Better assertion with less wait time

lib/yamatanooroti/vterm.rb

@@ -2,9 +2,11 @@
 require 'vterm'
 require 'pty'
 require 'io/console'
+require 'io/wait'
 
 module Yamatanooroti::VTermTestCaseModule
-  def start_terminal(height, width, command, wait: 0.1, startup_message: nil)
+  def start_terminal(height, width, command, wait: 0.01, timeout: 2, startup_message: nil)
+    @timeout = timeout
     @wait = wait
     @result = nil
 
@@ -18,14 +20,10 @@ def start_terminal(height, width, command, wait: 0.1, startup_message: nil)
 
     case startup_message
     when String
-      @startup_message = ->(message) { message.start_with?(startup_message) }
+      wait_startup_message { |message| message.start_with?(startup_message) }
     when Regexp
-      @startup_message = ->(message) { startup_message.match?(message) }
-    else
-      @startup_message = nil
+      wait_startup_message { |message| startup_message.match?(message) }
     end
-
-    sync
   end
 
   def write(str)
@@ -42,76 +40,126 @@ def write(str)
       end
     end
     @pty_input.write(str_to_write)
-    sync
+    # Write str (e.g. `exit`) to pty_input might terminate the process.
+    try_sync
   end
 
   def close
-    sync
-    @pty_input.close
-    sync
-    Process.kill('KILL', @pid)
-    Process.waitpid(@pid)
+    begin
+      sync
+      @pty_input.close
+      sync
+    rescue IOError, Errno::EIO
+    end
+    begin
+      Process.kill('KILL', @pid)
+      Process.waitpid(@pid)
+    rescue Errno::ESRCH
+    end
   end
 
-  private def sync
-    startup_message = +'' if @startup_message
+  private def wait_startup_message
+    wait_until = Time.now + @timeout
+    chunks = +''
     loop do
-      sleep @wait
-      chunk = @pty_output.read_nonblock(1024)
-      if @startup_message
-        startup_message << chunk
-        if @startup_message.(startup_message)
-          @startup_message = nil
-          chunk = startup_message
-        else
-          redo
-        end
+      wait = wait_until - Time.now
+      if wait.negative? || !@pty_output.wait_readable(wait)
+        raise "Startup message didn't arrive within timeout: #{chunks.inspect}"
       end
-      @vterm.write(chunk)
-      chunk = @vterm.read
-      @pty_input.write(chunk)
-    rescue Errno::EAGAIN, Errno::EWOULDBLOCK
-      retry if @startup_message
-      break
-    rescue Errno::EIO # EOF
-      retry if @startup_message
-      break
-    rescue IO::EAGAINWaitReadable # emtpy buffer
-      retry if @startup_message
-      break
+
+      chunk = @pty_output.read_nonblock(65536)
+      vterm_write(chunk)
+      chunks << chunk
+      break if yield chunks
+    end
+  end
+
+  private def vterm_write(chunk)
+    @vterm.write(chunk)
+    response = @vterm.read
+    begin
+      @pty_input.write(response)
+    rescue Errno::EIO
+      # In case process terminates suddenly after writing "\e[6n"
+    end
+    @result = nil
+  end
+
+  private def sync(wait = @wait)
+    sync_until = Time.now + @timeout
+    while @pty_output.wait_readable(wait)
+      vterm_write(@pty_output.read_nonblock(65536))
+      break if Time.now > sync_until
     end
   end
 
+  private def try_sync(wait = @wait)
+    sync(wait)
+    true
+  rescue IOError, Errno::EIO
+    false
+  end
+
+
   def result
-    return @result if @result
-    @result = []
+    try_sync(0)
+    @result ||= retrieve_screen
+  end
+
+  private def retrieve_screen
+    result = []
     rows, cols = @vterm.size
     rows.times do |r|
-      @result << +''
+      result << +''
       cols.times do |c|
         cell = @screen.cell_at(r, c)
         if cell.char # The second cell of fullwidth char will be nil.
           if cell.char.empty?
             # There will be no char to the left of the rendered area if moves
             # the cursor.
-            @result.last << ' '
+            result.last << ' '
           else
-            @result.last << cell.char
+            result.last << cell.char
           end
         end
       end
-      @result.last.gsub!(/ *$/, '')
+      result.last.gsub!(/ *$/, '')
+    end
+    result
+  end
+
+  private def retryable_screen_assertion_with_proc(check_proc, assert_proc, convert_proc = :itself.to_proc)
+    retry_until = Time.now + @timeout
+    while Time.now < retry_until
+      break unless try_sync
+
+      @result ||= retrieve_screen
+      break if check_proc.call(convert_proc.call(@result))
     end
-    @result
+    @result ||= retrieve_screen
+    assert_proc.call(convert_proc.call(@result))
   end
 
   def assert_screen(expected_lines, message = nil)
-    actual_lines = result
+    lines_to_string = ->(lines) { lines.join("\n").sub(/\n*\z/, "\n") }
     case expected_lines
     when Array
-      assert_equal(expected_lines, actual_lines, message)
+      retryable_screen_assertion_with_proc(
+        ->(actual) { expected_lines == actual },
+        ->(actual) { assert_equal(expected_lines, actual, message) }
+      )
     when String
-      assert_equal(expected_lines, actual_lines.join("\n").sub(/\n*\z/, "\n"), message)
+      retryable_screen_assertion_with_proc(
+        ->(actual) { expected_lines == actual },
+        ->(actual) { assert_equal(expected_lines, actual, message) },
+        lines_to_string
+      )
+    when Regexp
+      retryable_screen_assertion_with_proc(
+        ->(actual) { expected_lines.match?(actual) },
+        ->(actual) { assert_match(expected_lines, actual, message) },
+        lines_to_string
+      )
     end
   end
 end

lib/yamatanooroti/windows.rb

@@ -516,15 +516,19 @@ def close
   end
 
   def result
-    @result
+    @result || retrieve_screen
   end
 
   def assert_screen(expected_lines, message = nil)
+    actual_lines = result
+    actual_string = actual_lines.join("\n").sub(/\n*\z/, "\n")
     case expected_lines
     when Array
-      assert_equal(expected_lines, @result, message)
+      assert_equal(expected_lines, actual_lines, message)
     when String
-      assert_equal(expected_lines, @result.join("\n").sub(/\n*\z/, "\n"), message)
+      assert_equal(expected_lines, actual_string, message)
+    when Regexp
+      assert_match(expected_lines, actual_string, message)
     end
   end
 

test/yamatanooroti/test_multiplatform.rb

@@ -2,8 +2,7 @@
 
 class Yamatanooroti::TestMultiplatform < Yamatanooroti::TestCase
   def setup
-    start_terminal(5, 30, ['ruby', 'bin/simple_repl'])
-    sleep 0.5
+    start_terminal(5, 30, ['ruby', 'bin/simple_repl'], startup_message: 'prompt>')
   end
 
   def test_example
@@ -17,10 +16,29 @@ def test_example
     EOC
   end
 
-  def test_result
+  def test_result_repeatedly
     write(":a\n")
-    close
+    assert_screen(/=> :a\nprompt>/)
     assert_equal(['prompt> :a', '=> :a', 'prompt>', '', ''], result)
+    write(":b\n")
+    assert_screen(/=> :b\nprompt>/)
+    assert_equal(['prompt> :a', '=> :a', 'prompt> :b', '=> :b', 'prompt>'], result)
+    close
+  end
+
+  def test_assert_screen_retries
+    write("sleep 1\n")
+    assert_screen(/=> 1\nprompt>/)
+    assert_equal(['prompt> sleep 1', '=> 1', 'prompt>', '', ''], result)
+    close
+  end
+
+  def test_assert_screen_timeout
+    write("sleep 3\n")
+    assert_raise do
+      assert_screen(/=> 3\nprompt>/)
+    end
+    close
   end
 
   def test_auto_wrap
@@ -38,12 +56,14 @@ def test_auto_wrap
   def test_fullwidth
     write(":あ\n")
     close
+    assert_screen(/=> :あ\nprompt>/)
     assert_equal(['prompt> :あ', '=> :あ', 'prompt>', '', ''], result)
   end
 
   def test_two_fullwidth
     write(":あい\n")
     close
+    assert_screen(/=> :あい\nprompt>/)
     assert_equal(['prompt> :あい', '=> :あい', 'prompt>', '', ''], result)
   end
 end

test/yamatanooroti/test_run_ruby.rb

@@ -26,19 +26,20 @@ def test_wait_for_startup_message
 
   def test_move_cursor_and_render
     start_terminal(5, 30, ['ruby', '-rio/console', '-e', 'STDOUT.puts(?A);STDOUT.goto(2,2);STDOUT.puts(?B)'])
+    assert_screen(/^  B/)
     close
     assert_equal(['A', '', '  B', '', ''], result)
   end
 
   def test_meta_key
     get_into_tmpdir
-    start_terminal(5, 30, ['ruby', '-rreline', '-e', 'Reline.readline(%{?})'])
+    start_terminal(5, 30, ['ruby', '-rreline', '-e', 'Reline.readline(%{>>>})'], startup_message: />{3}/)
     write('aaa ccc')
     write("\M-b")
     write('bbb ')
     close
     assert_screen(<<~EOC)
-      ?aaa bbb ccc
+      >>>aaa bbb ccc
     EOC
   ensure
     get_out_from_tmpdir