Use Erlang logger handlers in ExUnit.CaptureLog (#12408)

Author: whatyouhide

lib/ex_unit/lib/ex_unit/capture_log.ex

@@ -37,9 +37,7 @@ defmodule ExUnit.CaptureLog do
 
   """
 
-  alias Logger.Backends.Console
-
-  @compile {:no_warn_undefined, [Logger, Logger.Backends.Internal]}
+  @compile {:no_warn_undefined, Logger}
 
   @doc """
   Captures Logger messages generated when evaluating `fun`.
@@ -93,6 +91,7 @@ defmodule ExUnit.CaptureLog do
 
       assert result == 4
       assert log =~ "log msg"
+
   """
   @doc since: "1.13.0"
   @spec with_log(keyword, (-> result)) :: {result, String.t()} when result: any
@@ -101,74 +100,22 @@ defmodule ExUnit.CaptureLog do
     {:ok, string_io} = StringIO.open("")
 
     try do
-      :ok = add_capture(string_io, opts)
-      ref = ExUnit.CaptureServer.log_capture_on(self())
+      ref = ExUnit.CaptureServer.log_capture_on(self(), string_io, opts)
 
       try do
         fun.()
       after
         :ok = Logger.flush()
         :ok = ExUnit.CaptureServer.log_capture_off(ref)
-        :ok = remove_capture(string_io)
       end
     catch
       kind, reason ->
         _ = StringIO.close(string_io)
         :erlang.raise(kind, reason, __STACKTRACE__)
     else
       result ->
-        {:ok, content} = StringIO.close(string_io)
-        {result, elem(content, 1)}
-    end
-  end
-
-  defp add_capture(pid, opts) do
-    case :proc_lib.start(__MODULE__, :init_proxy, [pid, opts, self()]) do
-      :ok ->
-        :ok
-
-      :noproc ->
-        raise "cannot capture_log/2 because the :logger application was not started"
-
-      {:error, reason} ->
-        mfa = {ExUnit.CaptureLog, :add_capture, [pid, opts]}
-        exit({reason, mfa})
-    end
-  end
-
-  @doc false
-  def init_proxy(pid, opts, parent) do
-    Logger.Backends.Internal.configure([])
-  catch
-    _, _ -> :proc_lib.init_ack(:noproc)
-  else
-    _ ->
-      case :gen_event.add_sup_handler(Logger, {Console, pid}, {Console, [device: pid] ++ opts}) do
-        :ok ->
-          ref = Process.monitor(parent)
-          :proc_lib.init_ack(:ok)
-
-          receive do
-            {:DOWN, ^ref, :process, ^parent, _reason} -> :ok
-            {:gen_event_EXIT, {Console, ^pid}, _reason} -> :ok
-          end
-
-        {:EXIT, reason} ->
-          :proc_lib.init_ack({:error, reason})
-
-        {:error, reason} ->
-          :proc_lib.init_ack({:error, reason})
-      end
-  end
-
-  defp remove_capture(pid) do
-    case :gen_event.delete_handler(Logger, {Console, pid}, :ok) do
-      :ok ->
-        :ok
-
-      {:error, :module_not_found} = error ->
-        mfa = {ExUnit.CaptureLog, :remove_capture, [pid]}
-        exit({error, mfa})
+        {:ok, {_input, output}} = StringIO.close(string_io)
+        {result, output}
     end
   end
 end

lib/ex_unit/lib/ex_unit/capture_server.ex

@@ -3,6 +3,7 @@ defmodule ExUnit.CaptureServer do
   @compile {:no_warn_undefined, Logger}
   @timeout :infinity
   @name __MODULE__
+  @ets __MODULE__
 
   use GenServer
 
@@ -22,8 +23,8 @@ defmodule ExUnit.CaptureServer do
     GenServer.call(@name, {:device_capture_off, ref}, @timeout)
   end
 
-  def log_capture_on(pid) do
-    GenServer.call(@name, {:log_capture_on, pid}, @timeout)
+  def log_capture_on(pid, string_io, opts) do
+    GenServer.call(@name, {:log_capture_on, pid, string_io, opts}, @timeout)
   end
 
   def log_capture_off(ref) do
@@ -32,7 +33,10 @@ defmodule ExUnit.CaptureServer do
 
   ## Callbacks
 
+  @impl true
   def init(:ok) do
+    :ets.new(@name, [:named_table, :public, :set])
+
     state = %{
       devices: %{},
       log_captures: %{},
@@ -42,6 +46,9 @@ defmodule ExUnit.CaptureServer do
     {:ok, state}
   end
 
+  @impl true
+  def handle_call(call, from, state)
+
   def handle_call({:device_capture_on, name, encoding, input}, {caller, _}, config) do
     capture_device(name, encoding, input, config, caller)
   end
@@ -59,11 +66,17 @@ defmodule ExUnit.CaptureServer do
     {:reply, :ok, release_device(ref, config)}
   end
 
-  def handle_call({:log_capture_on, pid}, _from, config) do
+  def handle_call({:log_capture_on, pid, string_io, opts}, _from, config) do
     ref = Process.monitor(pid)
     refs = Map.put(config.log_captures, ref, true)
 
+    {level, opts} = Keyword.pop(opts, :level)
+    true = :ets.insert(@ets, {ref, string_io, level || :all})
+
     if map_size(refs) == 1 do
+      formatter = Logger.default_formatter(opts)
+      :ok = :logger.add_handler(@name, __MODULE__, %{formatter: formatter})
+
       status =
         with {:ok, config} <- :logger.get_handler_config(:default),
              :ok <- :logger.remove_handler(:default) do
@@ -84,6 +97,7 @@ defmodule ExUnit.CaptureServer do
     {:reply, :ok, config}
   end
 
+  @impl true
   def handle_info({:DOWN, ref, _, _, _}, config) do
     config = remove_log_capture(ref, config)
     config = release_device(ref, config)
@@ -186,21 +200,39 @@ defmodule ExUnit.CaptureServer do
   end
 
   defp remove_log_capture(ref, %{log_captures: refs} = config) do
+    true = :ets.delete(@ets, ref)
+
     case Map.pop(refs, ref, false) do
       {true, refs} ->
-        maybe_add_console(refs, config.log_status)
+        maybe_revert_to_default_handler(refs, config.log_status)
         %{config | log_captures: refs}
 
       {false, _refs} ->
         config
     end
   end
 
-  defp maybe_add_console(refs, {:ok, %{module: module} = config}) when map_size(refs) == 0 do
-    :logger.add_handler(:default, module, config)
+  defp maybe_revert_to_default_handler(refs, status) when map_size(refs) == 0 do
+    :logger.remove_handler(@name)
+
+    with {:ok, %{module: module} = config} <- status do
+      :logger.add_handler(:default, module, config)
+    end
   end
 
-  defp maybe_add_console(_refs, _config) do
+  defp maybe_revert_to_default_handler(_refs, _config) do
     :ok
   end
+
+  ## :logger handler callback.
+
+  def log(event, %{} = config) do
+    %{formatter: {formatter_mod, formatter_config}} = config
+    chardata = formatter_mod.format(event, formatter_config)
+
+    for [string_io, level] <- :ets.match(@ets, {:_, :"$1", :"$2"}),
+        :logger.compare_levels(event.level, level) in [:gt, :eq] do
+      :ok = IO.write(string_io, chardata)
+    end
+  end
 end

lib/ex_unit/test/ex_unit/capture_log_test.exs

@@ -78,24 +78,39 @@ defmodule ExUnit.CaptureLogTest do
     end
   end
 
-  test "with_log" do
-    {4, log} =
-      with_log(fn ->
-        Logger.error("calculating...")
-        2 + 2
-      end)
+  describe "with_log/2" do
+    test "returns the result and the log" do
+      {result, log} =
+        with_log(fn ->
+          Logger.error("calculating...")
+          2 + 2
+        end)
+
+      assert result == 4
+      assert log =~ "calculating..."
+    end
+
+    test "respects the :format, :metadata, and :colors options" do
+      options = [format: "$metadata| $message", metadata: [:id], colors: [enabled: false]]
+
+      assert {4, log} =
+               with_log(options, fn ->
+                 Logger.info("hello", id: 123)
+                 2 + 2
+               end)
 
-    assert log =~ "calculating..."
+      assert log == "id=123 | hello"
+    end
   end
 
   defp wait_capture_removal() do
-    case :gen_event.which_handlers(Logger) do
-      [Logger.Backends.Config] ->
-        :ok
-
-      _otherwise ->
+    case Enum.map(:logger.get_handler_config(), & &1.id) do
+      [ExUnit.CaptureServer] ->
         Process.sleep(20)
         wait_capture_removal()
+
+      [:default] ->
+        :ok
     end
   end
 end