Support IO capture of group leader from external process (#13374)

Author: voughtdq

lib/ex_unit/lib/ex_unit/capture_io.ex

@@ -29,15 +29,15 @@ defmodule ExUnit.CaptureIO do
 
   Returns the binary which is the captured output.
 
-  By default, `capture_io` replaces the `group_leader` (`:stdio`)
-  for the current process. Capturing the group leader is done per
-  process and therefore can be done concurrently.
+  By default, `capture_io` replaces the `Process.group_leader/0` of the current
+  process, which is the process used by default for all IO operations. Capturing
+  the group leader of the current process is safe to run concurrently, under
+  `async: true` tests. You may also explicitly capture the group leader of
+  another process, however that is not safe to do concurrently.
 
-  However, the capturing of any other named device, such as `:stderr`,
-  happens globally and persists until the function has ended. While this means
-  it is safe to run your tests with `async: true` in many cases, captured output
-  may include output from a different test and care must be taken when using
-  `capture_io` with a named process asynchronously.
+  You may also capture any other named IO device, such as `:stderr`. This is
+  also safe to run concurrently but, if several tests are writting to the same
+  device at once, captured output may include output from a different test.
 
   A developer can set a string as an input. The default input is an empty
   string. If capturing a named device asynchronously, an input can only be given
@@ -51,15 +51,28 @@ defmodule ExUnit.CaptureIO do
 
   ## IO devices
 
-  You may capture the IO from any registered IO device. The device name given
-  must be an atom representing the name of a registered process. In addition,
-  Elixir provides two shortcuts:
+  You may capture the IO of the group leader of any process, by passing a `pid`
+  as argument, or from any registered IO device given as an `atom`. Here are
+  some example values:
 
-    * `:stdio` - a shortcut for `:standard_io`, which maps to
-      the current `Process.group_leader/0` in Erlang
+    * `:stdio`, `:standard_io` - a shortcut for capturing the group leader
+      of the current process. It is equivalent to passing `self()` as the
+      first argument. This is safe to run concurrently and captures only
+      the of the current process or any child process spawned inside the
+      given function
 
-    * `:stderr` - a shortcut for the named process `:standard_error`
-      provided in Erlang
+    * `:stderr`, `:standard_error` - captures all IO to standard error
+      (represented internally by an Erlang process named `:standard_error`).
+      This is safe to run concurrently but it will capture the output
+      of any other test writing to the same named device
+
+    * any other atom - captures all IO to the given device given by the
+      atom. This is safe to run concurrently but it will capture the output
+      of any other test writing to the same named device
+
+    * any other pid (since v1.17.0) - captures all IO to the group leader
+      of the given process. This option is not safe to run concurrently
+      if the pid is not `self()`. Tests using this value must set `async: true`
 
   ## Options
 
@@ -91,10 +104,10 @@ defmodule ExUnit.CaptureIO do
       ...> end) == "this is input"
       true
 
-  Note it is fine to use `==` with standard IO, because the content is captured
-  per test process. However, `:stderr` is shared across all tests, so you will
-  want to use `=~` instead of `==` for assertions on `:stderr` if your tests
-  are async:
+  Note it is fine to use `==` with `:stdio` (the default IO device), because
+  the content is captured per test process. However, `:stderr` is shared
+  across all tests, so you will want to use `=~` instead of `==` for assertions
+  on `:stderr` if your tests are async:
 
       iex> capture_io(:stderr, fn -> IO.write(:stderr, "john") end) =~ "john"
       true
@@ -110,6 +123,14 @@ defmodule ExUnit.CaptureIO do
   Otherwise, if the standard error of any other test is captured, the test will
   fail.
 
+  To capture the IO from another process, you can pass a `pid`:
+
+      capture_io(GenServer.whereis(MyServer), fn ->
+        GenServer.call(MyServer, :do_something)
+      end)
+
+  Tests that directly capture a PID cannot run concurrently.
+
   ## Returning values
 
   As seen in the examples above, `capture_io` returns the captured output.
@@ -127,14 +148,19 @@ defmodule ExUnit.CaptureIO do
 
   See `capture_io/1` for more information.
   """
-  @spec capture_io(atom() | String.t() | keyword(), (-> any())) :: String.t()
-  def capture_io(device_input_or_options, fun)
+  @spec capture_io(atom() | pid() | String.t() | keyword(), (-> any())) :: String.t()
+  def capture_io(device_pid_input_or_options, fun)
 
   def capture_io(device, fun) when is_atom(device) and is_function(fun, 0) do
     {_result, capture} = with_io(device, fun)
     capture
   end
 
+  def capture_io(pid, fun) when is_pid(pid) and is_function(fun, 0) do
+    {_result, capture} = with_io(pid, fun)
+    capture
+  end
+
   def capture_io(input, fun) when is_binary(input) and is_function(fun, 0) do
     {_result, capture} = with_io(input, fun)
     capture
@@ -150,8 +176,8 @@ defmodule ExUnit.CaptureIO do
 
   See `capture_io/1` for more information.
   """
-  @spec capture_io(atom(), String.t() | keyword(), (-> any())) :: String.t()
-  def capture_io(device, input_or_options, fun)
+  @spec capture_io(atom() | pid(), String.t() | keyword(), (-> any())) :: String.t()
+  def capture_io(device_or_pid, input_or_options, fun)
 
   def capture_io(device, input, fun)
       when is_atom(device) and is_binary(input) and is_function(fun, 0) do
@@ -165,6 +191,18 @@ defmodule ExUnit.CaptureIO do
     capture
   end
 
+  def capture_io(pid, input, fun)
+      when is_pid(pid) and is_binary(input) and is_function(fun, 0) do
+    {_result, capture} = with_io(pid, input, fun)
+    capture
+  end
+
+  def capture_io(pid, options, fun)
+      when is_pid(pid) and is_list(options) and is_function(fun, 0) do
+    {_result, capture} = with_io(pid, options, fun)
+    capture
+  end
+
   @doc ~S"""
   Invokes the given `fun` and returns the result and captured output.
 
@@ -194,13 +232,17 @@ defmodule ExUnit.CaptureIO do
   See `with_io/1` for more information.
   """
   @doc since: "1.13.0"
-  @spec with_io(atom() | String.t() | keyword(), (-> any())) :: {any(), String.t()}
-  def with_io(device_input_or_options, fun)
+  @spec with_io(atom() | pid() | String.t() | keyword(), (-> any())) :: {any(), String.t()}
+  def with_io(device_pid_input_or_options, fun)
 
   def with_io(device, fun) when is_atom(device) and is_function(fun, 0) do
     with_io(device, [], fun)
   end
 
+  def with_io(pid, fun) when is_pid(pid) and is_function(fun, 0) do
+    with_io(pid, [], fun)
+  end
+
   def with_io(input, fun) when is_binary(input) and is_function(fun, 0) do
     with_io(:stdio, [input: input], fun)
   end
@@ -215,8 +257,8 @@ defmodule ExUnit.CaptureIO do
   See `with_io/1` for more information.
   """
   @doc since: "1.13.0"
-  @spec with_io(atom(), String.t() | keyword(), (-> any())) :: {any(), String.t()}
-  def with_io(device, input_or_options, fun)
+  @spec with_io(atom() | pid(), String.t() | keyword(), (-> any())) :: {any(), String.t()}
+  def with_io(device_or_pid, input_or_options, fun)
 
   def with_io(device, input, fun)
       when is_atom(device) and is_binary(input) and is_function(fun, 0) do
@@ -228,23 +270,35 @@ defmodule ExUnit.CaptureIO do
     do_with_io(map_dev(device), options, fun)
   end
 
+  def with_io(pid, input, fun)
+      when is_pid(pid) and is_binary(input) and is_function(fun, 0) do
+    with_io(pid, [input: input], fun)
+  end
+
+  def with_io(pid, options, fun)
+      when is_pid(pid) and is_list(options) and is_function(fun, 0) do
+    do_with_io(pid, options, fun)
+  end
+
   defp map_dev(:stdio), do: :standard_io
   defp map_dev(:stderr), do: :standard_error
   defp map_dev(other), do: other
 
-  defp do_with_io(:standard_io, options, fun) do
+  defp do_with_io(device_or_pid, options, fun)
+       when device_or_pid == :standard_io or is_pid(device_or_pid) do
     prompt_config = Keyword.get(options, :capture_prompt, true)
     encoding = Keyword.get(options, :encoding, :unicode)
     input = Keyword.get(options, :input, "")
 
     original_gl = Process.group_leader()
     {:ok, capture_gl} = StringIO.open(input, capture_prompt: prompt_config, encoding: encoding)
+    pid = if is_pid(device_or_pid), do: device_or_pid, else: self()
 
     try do
-      Process.group_leader(self(), capture_gl)
+      Process.group_leader(pid, capture_gl)
       do_capture_gl(capture_gl, fun)
     after
-      Process.group_leader(self(), original_gl)
+      Process.group_leader(pid, original_gl)
     end
   end
 

lib/ex_unit/test/ex_unit/capture_io_test.exs

@@ -28,6 +28,38 @@ defmodule ExUnit.CaptureIOTest do
     end
   end
 
+  defmodule MockProc do
+    use GenServer
+
+    def start_link do
+      GenServer.start_link(__MODULE__, [])
+    end
+
+    @impl GenServer
+    def init(_), do: {:ok, nil}
+
+    @impl GenServer
+    def handle_call({:stdio, message}, _from, state) do
+      IO.puts(message)
+      {:reply, :ok, state}
+    end
+
+    @impl GenServer
+    def handle_call({:prompt, prompt}, _from, state) do
+      prompt
+      |> IO.gets()
+      |> IO.puts()
+
+      {:reply, :ok, state}
+    end
+
+    @impl GenServer
+    def handle_call({:stderr, message}, _from, state) do
+      IO.puts(:stderr, message)
+      {:reply, :ok, state}
+    end
+  end
+
   import ExUnit.CaptureIO
   doctest ExUnit.CaptureIO, import: true
 
@@ -469,6 +501,32 @@ defmodule ExUnit.CaptureIOTest do
     end
   end
 
+  test "capture_io with a separate process" do
+    {:ok, pid} = MockProc.start_link()
+
+    assert capture_io(pid, fn ->
+             GenServer.call(pid, {:stdio, "a"})
+           end) == "a\n"
+
+    assert capture_io(pid, [input: "b"], fn ->
+             GenServer.call(pid, {:prompt, "> "})
+           end) == "> b\n"
+
+    assert capture_io(pid, "c", fn ->
+             GenServer.call(pid, {:prompt, "> "})
+           end) == "> c\n"
+
+    assert capture_io(pid, [input: "d", capture_prompt: false], fn ->
+             GenServer.call(pid, {:prompt, "> "})
+           end) == "d\n"
+
+    assert capture_io(:stderr, fn ->
+             GenServer.call(pid, {:stderr, "uhoh"})
+           end) == "uhoh\n"
+
+    GenServer.stop(pid)
+  end
+
   test "with_io" do
     assert with_io(fn ->
              :io.put_chars("xyz")