ex_unit: Add :capture_io tag

lib/elixir/test/elixir/kernel/cli_test.exs

@@ -39,29 +39,28 @@ end
 defmodule Kernel.CLITest do
   use ExUnit.Case, async: true
 
-  import ExUnit.CaptureIO
-
   defp run(argv) do
     {config, argv} = Kernel.CLI.parse_argv(Enum.map(argv, &String.to_charlist/1))
     assert Kernel.CLI.process_commands(config) == []
     Enum.map(argv, &IO.chardata_to_string/1)
   end
 
-  test "argv handling" do
-    assert capture_io(fn ->
-             assert run(["-e", "IO.puts :ok", "sample.exs", "-o", "1", "2"]) ==
-                      ["sample.exs", "-o", "1", "2"]
-           end) == "ok\n"
-
-    assert capture_io(fn ->
-             assert run(["-e", "IO.puts :ok", "--", "sample.exs", "-o", "1", "2"]) ==
-                      ["sample.exs", "-o", "1", "2"]
-           end) == "ok\n"
-
-    assert capture_io(fn ->
-             assert run(["-e", "", "--", "sample.exs", "-o", "1", "2"]) ==
-                      ["sample.exs", "-o", "1", "2"]
-           end)
+  @tag :capture_io
+  test "argv handling", %{capture_io: io} do
+    assert run(["-e", "IO.puts :ok1", "sample.exs", "-o", "1", "2"]) ==
+             ["sample.exs", "-o", "1", "2"]
+
+    assert StringIO.flush(io) == "ok1\n"
+
+    assert run(["-e", "IO.puts :ok2", "--", "sample.exs", "-o", "1", "2"]) ==
+             ["sample.exs", "-o", "1", "2"]
+
+    assert StringIO.flush(io) == "ok2\n"
+
+    assert run(["-e", "", "--", "sample.exs", "-o", "1", "2"]) ==
+             ["sample.exs", "-o", "1", "2"]
+
+    assert StringIO.flush(io) == ""
   end
 end
 

lib/ex_unit/lib/ex_unit.ex

@@ -149,10 +149,21 @@ defmodule ExUnit do
       * `:time` - the duration in microseconds of the test's runtime
       * `:tags` - the test tags
       * `:logs` - the captured logs
+      * `:capture_io` - (since v1.20.0) the captured IO
       * `:parameters` - the test parameters
 
     """
-    defstruct [:name, :case, :module, :state, time: 0, tags: %{}, logs: "", parameters: %{}]
+    defstruct [
+      :name,
+      :case,
+      :module,
+      :state,
+      time: 0,
+      tags: %{},
+      logs: "",
+      capture_io: "",
+      parameters: %{}
+    ]
 
     # TODO: Remove the `:case` field on v2.0
     @type t :: %__MODULE__{

lib/ex_unit/lib/ex_unit/case.ex

@@ -183,6 +183,8 @@ defmodule ExUnit.Case do
 
   The following tags customize how tests behave:
 
+    * `:capture_io` - (since v1.20.0) see the "IO Capture" section below
+
     * `:capture_log` - see the "Log Capture" section below
 
     * `:skip` - skips the test with the given reason
@@ -258,6 +260,34 @@ defmodule ExUnit.Case do
   Keep in mind that all tests are included by default, so unless they are
   excluded first, the `include` option has no effect.
 
+  ## IO Capture
+
+  ExUnit can optionally suppress printing of standard output messages generated
+  during a test. Messages generated while running a test are captured and
+  only if the test fails are they printed to aid with debugging.
+
+  The captured IO is available in the test context under `:capture_io`
+  key and can be read using `StringIO.flush/1`:
+
+      defmodule MyTest do
+        use ExUnit.Case, async: true
+
+        @tag :capture_io
+        test "with io", %{capture_io: io} do
+          IO.puts("Hello, World!")
+
+          assert StringIO.flush(io) == "Hello, World!\n"
+        end
+      end
+
+  As with other tags, `:capture_io` can also be set as `@moduletag` and
+  `@describetag`.
+
+  Since `setup_all` blocks don't belong to a specific test, standard output
+  messages generated in them (or between tests) are never captured.
+
+  See also `ExUnit.CaptureIO`.
+
   ## Log Capture
 
   ExUnit can optionally suppress printing of log messages that are generated
@@ -278,6 +308,8 @@ defmodule ExUnit.Case do
 
       config :logger, :default_handler, false
 
+  See also `ExUnit.CaptureLog`.
+
   ## Tmp Dir
 
   ExUnit automatically creates a temporary directory for tests tagged with

lib/ex_unit/lib/ex_unit/cli_formatter.ex

@@ -132,6 +132,7 @@ defmodule ExUnit.CLIFormatter do
       )
 
     print_failure(formatted, config)
+    print_capture_io(test.capture_io)
     print_logs(test.logs)
 
     test_counter = update_test_counter(config.test_counter, test)
@@ -519,4 +520,12 @@ defmodule ExUnit.CLIFormatter do
     output = String.replace(output, "\n", indent)
     IO.puts(["     The following output was logged:", indent | output])
   end
+
+  defp print_capture_io(""), do: nil
+
+  defp print_capture_io(output) do
+    indent = "\n     "
+    output = String.replace(output, "\n", indent)
+    IO.puts(["     The following standard output was captured:", indent | output])
+  end
 end

lib/ex_unit/lib/ex_unit/runner.ex

@@ -439,16 +439,19 @@ defmodule ExUnit.Runner do
       generate_test_seed(seed, test, rand_algorithm)
       context = context |> Map.merge(test.tags) |> Map.put(:test_pid, self())
       capture_log = Map.get(context, :capture_log, capture_log)
+      capture_io = Map.get(context, :capture_io, false)
 
       {time, test} =
         :timer.tc(
           maybe_capture_log(capture_log, test, fn ->
-            context = maybe_create_tmp_dir(context, test)
-
-            case exec_test_setup(test, context) do
-              {:ok, context} -> exec_test(test, context)
-              {:error, test} -> test
-            end
+            maybe_capture_io(capture_io, context, fn context ->
+              context = maybe_create_tmp_dir(context, test)
+
+              case exec_test_setup(test, context) do
+                {:ok, context} -> exec_test(test, context)
+                {:error, test} -> test
+              end
+            end)
           end)
         )
 
@@ -482,6 +485,32 @@ defmodule ExUnit.Runner do
     end
   end
 
+  defp maybe_capture_io(true, context, fun) do
+    {:ok, gl} = StringIO.open("")
+    Process.group_leader(self(), gl)
+    context = put_in(context.capture_io, gl)
+    test = fun.(context)
+    put_in(test.capture_io, StringIO.flush(gl))
+  end
+
+  defp maybe_capture_io(false, context, fun) do
+    fun.(context)
+  end
+
+  defp maybe_capture_io(other, _context, _fun) do
+    raise ArgumentError, """
+    invalid value for @tag :capture_io, expected one of:
+
+        @tag :capture_io
+        @tag capture_io: true
+        @tag capture_io: false
+
+    got:
+
+        @tag capture_io: #{inspect(other)}
+    """
+  end
+
   defp receive_test_reply(test, test_pid, test_ref, timeout) do
     receive do
       {^test_pid, :test_finished, test} ->

lib/ex_unit/test/ex_unit_test.exs

@@ -350,6 +350,43 @@ defmodule ExUnitTest do
     assert output =~ "\n1 test, 1 failure (3 excluded)\n"
   end
 
+  test "io capturing" do
+    defmodule IOCapturingTest do
+      use ExUnit.Case
+
+      @tag :capture_io
+      test "one" do
+        # test successful, captured "one" isn't printed
+        IO.puts("one")
+        assert 1 == 1
+      end
+
+      @tag :capture_io
+      test "two" do
+        # test failed, captured "two" is printed
+        IO.puts("two")
+        assert 1 == 2
+      end
+
+      @tag :capture_io
+      test "three, four", %{capture_io: io} do
+        # io is flushed, captured "three" isn't printed
+        IO.puts("three")
+        assert StringIO.flush(io) == "three\n"
+
+        # test failed, captured "four" is printed
+        IO.puts("four")
+        assert 1 == 2
+      end
+    end
+
+    output = capture_io(&ExUnit.run/0)
+    refute output =~ "one\n"
+    assert output =~ "two\n"
+    refute output =~ "three\n"
+    assert output =~ "four\n"
+  end
+
   test "log capturing" do
     defmodule LogCapturingTest do
       use ExUnit.Case