Support on_exit for registering callbacks in ExUnit.

Author: MSch

lib/ex_unit/lib/ex_unit/callbacks.ex

@@ -143,6 +143,14 @@ defmodule ExUnit.Callbacks do
     end
   end
 
+  @spec on_exit(term, (() -> term)) :: :ok
+  def on_exit(ref \\ make_ref, callback) do
+    send(Process.get(ExUnit.Runner.Pid), {self(), :register_on_exit_callback, ref, callback})
+    receive do
+      {:registered_on_exit_callback, ^ref} -> :ok
+    end
+  end
+
   ## Helpers
 
   @doc false

lib/ex_unit/lib/ex_unit/runner.ex

@@ -195,6 +195,7 @@ defmodule ExUnit.Runner do
 
     {test_pid, test_ref} =
       spawn_monitor(fn ->
+        Process.put(ExUnit.Runner.Pid, self_pid)
         {us, test} =
           :timer.tc(fn ->
             case exec_test_setup(test, context) do
@@ -211,10 +212,24 @@ defmodule ExUnit.Runner do
         send self_pid, {self, :test_finished, %{test | time: us}}
       end)
 
+    loop_test(test, test_pid, test_ref, HashDict.new)
+  end
+
+  defp loop_test(test, test_pid, test_ref, on_exit_callbacks) do
     receive do
+      {^test_pid, :register_on_exit_callback, ref, callback} ->
+        on_exit_callbacks = Dict.put(on_exit_callbacks, ref, callback)
+        send(test_pid, {:registered_on_exit_callback, ref})
+        loop_test(test, test_pid, test_ref, on_exit_callbacks)
       {^test_pid, :test_finished, test} ->
-        test
+        on_exit_state = exec_on_exit_callbacks(on_exit_callbacks)
+        if nil?(test.state) do
+          %{test | state: on_exit_state}
+        else
+          test
+        end
       {:DOWN, ^test_ref, :process, ^test_pid, error} ->
+        exec_on_exit_callbacks(on_exit_callbacks)
         %{test | state: {:failed, {{:EXIT, test_pid}, error, []}}}
     end
   end
@@ -249,6 +264,59 @@ defmodule ExUnit.Runner do
       end
   end
 
+  defp exec_on_exit_callbacks(callbacks) do
+    {runner_pid, runner_monitor, callback_state} = Enum.reduce Dict.values(callbacks), {nil, nil, nil}, fn(callback, {runner_pid, runner_monitor, callback_state}) ->
+      {runner_pid, runner_monitor} = ensure_alive_callback_runner(runner_pid, runner_monitor)
+      send(runner_pid, {:run, self(), callback})
+      receive do
+        {^runner_pid, nil} -> {runner_pid, runner_monitor, callback_state}
+        {^runner_pid, error} ->
+          if nil?(callback_state) do
+            {runner_pid, runner_monitor, error}
+          else
+            {runner_pid, runner_monitor, callback_state}
+          end
+        {:DOWN, ^runner_monitor, :process, ^runner_pid, error} ->
+          if nil?(callback_state) do
+            {nil, nil, {:failed, {{:EXIT, runner_pid}, error, []}}}
+          else
+            {nil, nil, callback_state}
+          end
+      end
+    end
+
+    if is_pid(runner_pid) and Process.alive?(runner_pid) do
+      send(runner_pid, :shutdown)
+      receive do
+        {:DOWN, ^runner_monitor, :process, ^runner_pid, _error} -> :ok
+      end
+    end
+
+    callback_state
+  end
+
+  defp ensure_alive_callback_runner(nil, nil), do: spawn_monitor(__MODULE__, :on_exit_runner_loop, [])
+  defp ensure_alive_callback_runner(runner_pid, runner_monitor), do: {runner_pid, runner_monitor}
+
+  def on_exit_runner_loop do
+    receive do
+      {:run, from, fun} ->
+        result = exec_on_exit_callback(fun)
+        send(from, {self(), result})
+        on_exit_runner_loop()
+      :shutdown -> :ok
+    end
+  end
+
+  defp exec_on_exit_callback(callback) do
+    callback.()
+    nil
+  catch
+    kind, error ->
+      {:failed, {kind, Exception.normalize(kind, error), pruned_stacktrace}}
+  end
+
+
   ## Helpers
 
   defp shuffle(%{seed: 0}, list) do

lib/ex_unit/test/ex_unit/callbacks_test.exs

@@ -132,6 +132,89 @@ defmodule ExUnit.CallbacksTest do
     assert capture_io(fn -> ExUnit.run end) =~
            "** (MatchError) no match of right hand side value: :error"
   end
+
+  test "doesn't choke on on_exit errors" do
+    defmodule OnExitErrorTest do
+      use ExUnit.Case, async: false
+
+      test "ok" do
+        on_exit fn ->
+          :ok = error
+        end
+
+        :ok
+      end
+
+      defp error, do: :error
+    end
+
+    ExUnit.configure(formatters: [ExUnit.CLIFormatter])
+
+    assert capture_io(fn -> ExUnit.run end) =~
+           "** (MatchError) no match of right hand side value: :error"
+  end
+
+  test "doesn't choke when on_exit exits" do
+    defmodule OnExitExitTest do
+      use ExUnit.Case, async: false
+
+      test "ok" do
+        on_exit fn ->
+          Process.exit(self(), :kill)
+        end
+
+        :ok
+      end
+    end
+
+    ExUnit.configure(formatters: [ExUnit.CLIFormatter])
+    assert capture_io(fn -> ExUnit.run end) =~
+           ">) killed"
+  end
+
+  test "runs multiple on_exit exits and overrides by ref" do
+    defmodule OnExitOverrideTest do
+      use ExUnit.Case, async: false
+
+      setup do
+        on_exit fn ->
+          IO.puts "on_exit setup run"
+        end
+
+        on_exit {:overridden, 1}, fn ->
+          IO.puts "on_exit 1 overridden -> not run"
+        end
+      end
+
+      test "ok" do
+        on_exit fn ->
+          IO.puts "simple on_exit run"
+        end
+
+        on_exit {:overridden, 2}, fn ->
+          IO.puts "on_exit 2 overridden -> not run"
+        end
+
+        on_exit {:overridden, 2}, fn ->
+          IO.puts "on_exit 2 overrides -> run"
+        end
+
+        on_exit {:overridden, 1}, fn ->
+          IO.puts "on_exit 1 overrides -> run"
+        end
+
+        :ok
+      end
+    end
+
+    ExUnit.configure(formatters: [ExUnit.CLIFormatter])
+    captured_id = capture_io(fn -> ExUnit.run end)
+    assert captured_id =~ "on_exit setup run"
+    assert captured_id =~ "simple on_exit run"
+    refute captured_id =~ "not run"
+    assert captured_id =~ "on_exit 1 overrides -> run"
+    assert captured_id =~ "on_exit 2 overrides -> run"
+  end
 end
 
 defmodule ExUnit.CallbacksNoTests do