From 0007821bb2f9a34b3c269660131c957bcc40eceb Mon Sep 17 00:00:00 2001
From: DanX <danx.exe@gmail.com>
Date: Fri, 11 Jan 2019 15:45:10 -0200
Subject: [PATCH] Add setup and teardown callbacks

---
 lib/que/job.ex        |  3 +++
 lib/que/worker.ex     | 26 ++++++++++++++++++++++-
 test/que/job_test.exs | 49 ++++++++++++++++++++++++++++++++++++++++---
 test/support.ex       | 19 +++++++++++++++++
 4 files changed, 93 insertions(+), 4 deletions(-)

diff --git a/lib/que/job.ex b/lib/que/job.ex
index e92e875..1ff42b0 100644
--- a/lib/que/job.ex
+++ b/lib/que/job.ex
@@ -59,6 +59,7 @@ defmodule Que.Job do
 
     {:ok, pid} =
       Que.Helpers.do_task(fn ->
+        job.worker.on_setup(job)
         job.worker.perform(job.arguments)
       end)
 
@@ -78,6 +79,7 @@ defmodule Que.Job do
 
     Que.Helpers.do_task(fn ->
       job.worker.on_success(job.arguments)
+      job.worker.on_teardown(job)
     end)
 
     %{ job | status: :completed, pid: nil, ref: nil }
@@ -96,6 +98,7 @@ defmodule Que.Job do
 
     Que.Helpers.do_task(fn ->
       job.worker.on_failure(job.arguments, error)
+      job.worker.on_teardown(job)
     end)
 
     %{ job | status: :failed, pid: nil, ref: nil }
diff --git a/lib/que/worker.ex b/lib/que/worker.ex
index 1bb4dc5..7f6876a 100644
--- a/lib/que/worker.ex
+++ b/lib/que/worker.ex
@@ -174,7 +174,15 @@ defmodule Que.Worker do
       end
 
 
-      defoverridable [on_success: 1, on_failure: 2]
+      def on_setup(_job) do
+      end
+
+
+      def on_teardown(_job) do
+      end
+
+
+      defoverridable [on_success: 1, on_failure: 2, on_setup: 1, on_teardown: 1]
 
 
 
@@ -232,4 +240,20 @@ defmodule Que.Worker do
   """
   @callback on_failure(arguments :: term, error :: tuple) :: term
 
+
+
+
+  @doc """
+  Optional callback that is executed before the job is started.
+  """
+  @callback on_setup(job :: term) :: term
+
+
+
+
+  @doc """
+  Optional callback that is executed after the job finishes,
+  both on success and failure.
+  """
+  @callback on_teardown(job :: term) :: term
 end
diff --git a/test/que/job_test.exs b/test/que/job_test.exs
index 7018c00..c39621e 100644
--- a/test/que/job_test.exs
+++ b/test/que/job_test.exs
@@ -6,6 +6,7 @@ defmodule Que.Test.Job do
   alias Que.Test.Meta.TestWorker
   alias Que.Test.Meta.SuccessWorker
   alias Que.Test.Meta.FailureWorker
+  alias Que.Test.Meta.SetupAndTeardownWorker
 
 
   test "#new builds a new Job struct with defaults" do
@@ -56,7 +57,7 @@ defmodule Que.Test.Job do
       refute job.pid    == nil
       refute job.ref    == nil
 
-      Helpers.wait
+      Helpers.wait_for_children
     end)
 
     assert capture =~ ~r/Starting/
@@ -64,6 +65,48 @@ defmodule Que.Test.Job do
   end
 
 
+  test "#perform triggers the worker's on_setup and on_teardown callbacks on success" do
+    capture = Helpers.capture_log(fn ->
+      job =
+        SetupAndTeardownWorker
+        |> Job.new
+        |> Job.perform
+        |> Job.handle_success
+
+      assert job.status == :completed
+      assert job.pid    == nil
+      assert job.ref    == nil
+
+      Helpers.wait_for_children
+    end)
+
+    assert capture =~ ~r/Completed/
+    assert capture =~ ~r/on_setup: %Que.Job/
+    assert capture =~ ~r/on_teardown: %Que.Job/
+  end
+
+
+  test "#perform triggers the worker's on_setup and on_teardown callbacks on failure" do
+    capture = Helpers.capture_log(fn ->
+      job =
+        SetupAndTeardownWorker
+        |> Job.new
+        |> Job.perform
+        |> Job.handle_failure("some error")
+
+      assert job.status == :failed
+      assert job.pid    == nil
+      assert job.ref    == nil
+
+      Helpers.wait_for_children
+    end)
+
+    assert capture =~ ~r/Failed/
+    assert capture =~ ~r/on_setup: %Que.Job/
+    assert capture =~ ~r/on_teardown: %Que.Job/
+  end
+
+
   test "#handle_success works as expected" do
     capture = Helpers.capture_log(fn ->
       job =
@@ -75,7 +118,7 @@ defmodule Que.Test.Job do
       assert job.pid    == nil
       assert job.ref    == nil
 
-      Helpers.wait
+      Helpers.wait_for_children
     end)
 
     assert capture =~ ~r/Completed/
@@ -94,7 +137,7 @@ defmodule Que.Test.Job do
       assert job.pid    == nil
       assert job.ref    == nil
 
-      Helpers.wait
+      Helpers.wait_for_children
     end)
 
     assert capture =~ ~r/Failed/
diff --git a/test/support.ex b/test/support.ex
index f81b896..53aa75b 100644
--- a/test/support.ex
+++ b/test/support.ex
@@ -56,6 +56,15 @@ defmodule Que.Test.Meta do
   end
 
 
+  defmodule SetupAndTeardownWorker do
+    use Que.Worker
+
+    def perform(args), do: Logger.debug("#{__MODULE__} - perform: #{inspect(args)}")
+    def on_setup(job), do: Logger.debug("#{__MODULE__} - on_setup: #{inspect(job)}")
+    def on_teardown(job), do: Logger.debug("#{__MODULE__} - on_teardown: #{inspect(job)}")
+  end
+
+
 
 
   # Helper Module for Tests
@@ -68,6 +77,16 @@ defmodule Que.Test.Meta do
       :timer.sleep(ms)
     end
 
+    def wait_for_children do
+      Task.Supervisor.children(Que.TaskSupervisor)
+      |> Enum.map(&Process.monitor/1)
+      |> Enum.each(fn ref ->
+        receive do
+          {:DOWN, ^ref, _, _, _} -> nil
+        end
+      end)
+    end
+
     # Captures IO output
     def capture_io(fun) do
       ExUnit.CaptureIO.capture_io(fun)