Merge pull request #3382 from lexmag/string_io-fixes

Fix `StringIO` processes leakage in the `ExUnit.CaptureIO`

Signed-off-by: José Valim <jose.valim@plataformatec.com.br>

Author: fishcakez

lib/elixir/test/elixir/string_io_test.exs

@@ -3,17 +3,12 @@ Code.require_file "test_helper.exs", __DIR__
 defmodule StringIOTest do
   use ExUnit.Case, async: true
 
-  test "start and stop" do
+  test "open and close" do
     {:ok, pid} = StringIO.open("")
     assert StringIO.close(pid) == {:ok, {"", ""}}
   end
 
-  test "start_link and stop" do
-    {:ok, pid} = StringIO.open("")
-    assert StringIO.close(pid) == {:ok, {"", ""}}
-  end
-
-  test "peek" do
+  test "contents" do
     {:ok, pid} = StringIO.open("abc")
     IO.write(pid, "edf")
     assert StringIO.contents(pid) == {"abc", "edf"}

lib/ex_unit/lib/ex_unit/capture_io.ex

@@ -9,13 +9,17 @@ defmodule ExUnit.CaptureIO do
 
         import ExUnit.CaptureIO
 
-        test :example do
-          assert capture_io(fn ->
-            IO.puts "a"
-          end) == "a\n"
+        test "checking the return value and the IO output" do
+          fun = fn ->
+            assert Enum.each(["some", "example"], &(IO.puts &1)) == :ok
+          end
+
+          assert capture_io(fun) == "some\nexample\n"
+
+          # Or use only: `capture_io(fun)` to silence the
+          # IO output (so only assert the return value)
         end
       end
-
   """
 
   @doc """
@@ -91,15 +95,13 @@ defmodule ExUnit.CaptureIO do
     prompt_config = Keyword.get(options, :capture_prompt, true)
     input = Keyword.get(options, :input, "")
 
-    original_gl = :erlang.group_leader
+    original_gl = Process.group_leader()
     {:ok, capture_gl} = StringIO.open(input, capture_prompt: prompt_config)
-    :erlang.group_leader(capture_gl, self)
-
     try do
-      fun.()
-      StringIO.close(capture_gl) |> elem(1) |> elem(1)
+      Process.group_leader(self(), capture_gl)
+      do_capture_io(capture_gl, fun)
     after
-      :erlang.group_leader(original_gl, self)
+      Process.group_leader(self(), original_gl)
     end
   end
 
@@ -119,8 +121,7 @@ defmodule ExUnit.CaptureIO do
     Process.register(capture_io, device)
 
     try do
-      fun.()
-      StringIO.close(capture_io) |> elem(1) |> elem(1)
+      do_capture_io(capture_io, fun)
     after
       try do
         Process.unregister(device)
@@ -131,4 +132,20 @@ defmodule ExUnit.CaptureIO do
       ExUnit.Server.remove_device(device)
     end
   end
+
+  defp do_capture_io(string_io, fun) do
+    try do
+       _ = fun.()
+      :ok
+    catch
+      kind, reason ->
+        stack = System.stacktrace()
+        _ = StringIO.close(string_io)
+        :erlang.raise(kind, reason, stack)
+    else
+      :ok ->
+        {:ok, output} = StringIO.close(string_io)
+        elem(output, 1)
+    end
+  end
 end

lib/ex_unit/test/ex_unit/capture_io_test.exs

@@ -28,6 +28,25 @@ defmodule ExUnit.CaptureIOTest do
   import ExUnit.CaptureIO
   doctest ExUnit.CaptureIO, import: true
 
+  test "no leakage on failures" do
+    group_leader = Process.group_leader()
+
+    test = self()
+    assert_raise ArgumentError, fn ->
+      capture_io(fn ->
+        send(test, {:string_io, Process.group_leader()})
+        raise ArgumentError
+      end)
+    end
+
+    receive do
+      {:string_io, pid} ->
+        ref = Process.monitor(pid)
+        assert_receive {:DOWN, ^ref, _, _, _}
+    end
+    assert Process.group_leader() == group_leader
+  end
+
   test "with no output" do
     assert capture_io(fn ->
     end) == ""
@@ -290,19 +309,14 @@ defmodule ExUnit.CaptureIOTest do
   end
 
   test "with assert inside" do
-    group_leader = :erlang.group_leader
-
     try do
       capture_io(fn ->
         assert false
       end)
     rescue
       error in [ExUnit.AssertionError] ->
-        "Expected truthy, got false" = error.message
+        assert error.message == "Expected truthy, got false"
     end
-
-    # Ensure no leakage on failures
-    assert group_leader == :erlang.group_leader
   end
 
   test "capture :stderr by two processes" do