capture_io captures a prompt when capturing of stdio

Author: Yuki Ito

lib/ex_unit/lib/ex_unit/capture_io.ex

@@ -28,6 +28,9 @@ defmodule ExUnit.CaptureIO do
   named device like `:stderr` is also possible globally by
   giving the registered device name explicitly as argument.
 
+  When capturing of `:stdio`, this function captures a prompt,
+  otherwise do not.
+
   A developer can set a string as an input. The default
   input is `:eof`.
 
@@ -40,9 +43,9 @@ defmodule ExUnit.CaptureIO do
       iex> capture_io(:stderr, fn -> IO.write(:stderr, "josé") end) == "josé"
       true
       iex> capture_io("this is input", fn->
-      ...>   input = IO.gets ""
+      ...>   input = IO.gets ">"
       ...>   IO.write input
-      ...> end) == "this is input"
+      ...> end) == ">this is input"
       true
 
   """
@@ -68,7 +71,7 @@ defmodule ExUnit.CaptureIO do
 
   defp do_capture_io(:standard_io, input, fun) do
     original_gl = :erlang.group_leader
-    capture_gl = new_group_leader(self, input)
+    capture_gl = new_group_leader(self, input, true)
     :erlang.group_leader(capture_gl, self)
 
     try do
@@ -105,12 +108,13 @@ defmodule ExUnit.CaptureIO do
     end
   end
 
-  defp new_group_leader(runner, input) do
-    spawn_link(fn -> group_leader_process(runner, input) end)
+  defp new_group_leader(runner, input, prompt_config // false) do
+    spawn_link(fn -> group_leader_process(runner, input, prompt_config) end)
   end
 
-  defp group_leader_process(runner, input) do
+  defp group_leader_process(runner, input, prompt_config) do
     register_input(input)
+    register_prompt_config(prompt_config)
     group_leader_loop(runner, :infinity, [])
   end
 
@@ -123,6 +127,10 @@ defmodule ExUnit.CaptureIO do
     set_input(chars)
   end
 
+  defp register_prompt_config(bool) do
+    Process.put(:capture_io_prompt_config, bool)
+  end
+
   defp set_input(:eof) do
     set_input([])
   end
@@ -135,6 +143,10 @@ defmodule ExUnit.CaptureIO do
     Process.get(:capture_io_input)
   end
 
+  defp need_prompt? do
+    Process.get(:capture_io_prompt_config)
+  end
+
   defp group_leader_loop(runner, wait, buf) do
     receive do
       { :io_request, from, reply_as, req } ->
@@ -181,28 +193,42 @@ defmodule ExUnit.CaptureIO do
     io_request({ :put_chars, mod, func, args }, buf)
   end
 
-  defp io_request({ :get_chars, _enc, _prompt, n }, buf) when n >= 0 do
-    { get_chars(n), buf }
+  defp io_request({ :get_chars, _enc, prompt, n }, buf) when n >= 0 do
+    io_request({ :get_chars, prompt, n }, buf)
   end
 
-  defp io_request({ :get_chars, _prompt, n }, buf) when n >= 0 do
+  defp io_request({ :get_chars, prompt, n }, buf) when n >= 0 do
+    if need_prompt? do
+      buf = [prompt|buf]
+    end
+
     { get_chars(n), buf }
   end
 
-  defp io_request({ :get_line, _prompt }, buf) do
-    { get_line, buf }
+  defp io_request({ :get_line, _enc, prompt }, buf) do
+    io_request({ :get_line, prompt }, buf)
   end
 
-  defp io_request({ :get_line, _enc, _prompt }, buf) do
+  defp io_request({ :get_line, prompt }, buf) do
+    if need_prompt? do
+      buf = [prompt|buf]
+    end
+
     { get_line, buf }
   end
 
-  defp io_request({ :get_until, _prompt, mod, fun, args }, buf) do
-    { get_until(mod, fun, args), buf }
+  defp io_request({ :get_until, _encoding, prompt, mod, fun, args}, buf) do
+    io_request({ :get_until, prompt, mod, fun, args}, buf)
   end
 
-  defp io_request({ :get_until, _encoding, _prompt, mod, fun, args}, buf) do
-    { get_until(mod, fun, args), buf }
+  defp io_request({ :get_until, prompt, mod, fun, args }, buf) do
+    { result, count } = get_until(mod, fun, args)
+
+    if need_prompt? do
+      buf = [:lists.duplicate(count, prompt)|buf]
+    end
+
+    { result, buf }
   end
 
   defp io_request({ :setopts, _opts }, buf) do
@@ -274,37 +300,37 @@ defmodule ExUnit.CaptureIO do
     do_get_until(input, mod, fun, args)
   end
 
-  defp do_get_until([], mod, fun, args, continuation // [])
+  defp do_get_until([], mod, fun, args, continuation // [], count // 0)
 
-  defp do_get_until([], mod, fun, args, continuation) do
+  defp do_get_until([], mod, fun, args, continuation, count) do
     case apply(mod, fun, [continuation, :eof | args]) do
       { :done, result, rest_chars } ->
         set_input(rest_chars)
-        result
+        { result, count + 1 }
       { :more, next_continuation } ->
-        do_get_until([], mod, fun, args, next_continuation)
+        do_get_until([], mod, fun, args, next_continuation, count + 1)
     end
   end
 
-  defp do_get_until(input, mod, fun, args, continuation) do
+  defp do_get_until(input, mod, fun, args, continuation, count) do
     { line, rest } = Enum.split_while(input, fn(char) -> char != ?\n end)
 
     case rest do
       [] ->
         case apply(mod, fun, [continuation, line | args]) do
           { :done, result, rest_chars } ->
             set_input(rest_chars)
-            result
+            { result, count + 1 }
           { :more, next_continuation } ->
-            do_get_until([], mod, fun, args, next_continuation)
+            do_get_until([], mod, fun, args, next_continuation, count + 1)
         end
       [_|t] ->
         case apply(mod, fun, [continuation, line ++ '\n' | args]) do
           { :done, result, rest_chars } ->
             set_input(rest_chars ++ t)
-            result
+            { result, count + 1 }
           { :more, next_continuation } ->
-            do_get_until(t, mod, fun, args, next_continuation)
+            do_get_until(t, mod, fun, args, next_continuation, count + 1)
         end
     end
   end

lib/ex_unit/test/ex_unit/capture_io_test.exs

@@ -78,7 +78,7 @@ defmodule ExUnit.CaptureIOTest do
   test :capture_io_with_get_chars do
     assert capture_io(fn ->
       :io.get_chars(">", 3)
-    end) == nil
+    end) == ">"
 
     capture_io(fn ->
       assert :io.get_chars(">", 3) == :eof
@@ -104,7 +104,7 @@ defmodule ExUnit.CaptureIOTest do
   test :capture_io_with_get_line do
     assert capture_io(fn ->
       :io.get_line ">"
-    end) == nil
+    end) == ">"
 
     capture_io(fn ->
       assert :io.get_line(">") == :eof
@@ -144,8 +144,16 @@ defmodule ExUnit.CaptureIOTest do
 
   test :capture_io_with_get_until do
     assert capture_io(fn ->
-      assert :io.scan_erl_form('>')
-    end) == nil
+      :io.scan_erl_form('>')
+    end) == ">"
+
+    assert capture_io("1.\n", fn ->
+      :io.scan_erl_form('>')
+    end) == ">"
+
+    assert capture_io("1\n.\n", fn ->
+       :io.scan_erl_form('>')
+    end) == ">>"
 
     capture_io(fn ->
       assert :io.scan_erl_form('>') == { :eof, 1 }