test: refactor utils

3 files changed, 78 insertions, 89 deletions

diff --git a/test/server_connection_test.rb b/test/server_connection_test.rb
index 39fc56e..1b72080 100644
--- a/test/server_connection_test.rb
+++ b/test/server_connection_test.rb
@@ -131,7 +131,7 @@ class ServerConnectionTest < Minitest::Test
 
     prepare.call {|con|
       con << Frame.new(type: :ping, flags: [], stream_id: 0, payload: "AAAAAAAA").assemble
-      last = sent_frames(con).last
+      last = sent_frames.last
       assert_equal(:ping, last.type)
       assert_equal([:ack], last.flags)
       assert_equal("AAAAAAAA", last.payload)
@@ -143,7 +143,7 @@ class ServerConnectionTest < Minitest::Test
     }
     prepare.call {|con|
       con << Frame.new(type: :ping, flags: [:ack], stream_id: 0, payload: "A" * 8).assemble
-      last = sent_frames(con).last
+      last = sent_frames.last
       refute_equal(:ping, last.type) if last
     }
   end
diff --git a/test/stream_test.rb b/test/stream_test.rb
index 90c0b7c..48c0f46 100644
--- a/test/stream_test.rb
+++ b/test/stream_test.rb
@@ -6,21 +6,11 @@ class StreamTest < Minitest::Test
   include ServerTestUtils
 
   def test_stream_reserve
-    prepare = -> &blk {
-      con = open_server_connection
-      stream = Stream.new(con, 2)
-      blk.call(stream)
-    }
-
-    prepare.call {|stream|
-      stream.instance_eval { @state = :idle }
-      refute_raises {
-        stream.reserve
-      }
+    open_new_stream {|stream|
+      stream.reserve
       assert_equal(:reserved_local, stream.state)
     }
-    prepare.call {|stream|
-      stream.instance_eval { @state = :open }
+    open_new_stream(:open) {|stream|
       assert_connection_error(:protocol_error) {
         stream.reserve
       }
@@ -28,34 +18,17 @@ class StreamTest < Minitest::Test
   end
 
   def test_stream_state_illegal_frame_type
-    test = -> (state, &blk) {
-      con = open_server_connection
-      stream = Stream.new(con, 2)
-      stream.instance_eval { @state = state }
-      blk.call(stream)
-    }
-
-    test.call(:idle) {|stream|
+    open_new_stream {|stream|
       assert_connection_error(:protocol_error) {
         stream.process_frame(Frame.new(type: :rst_stream, stream_id: stream.id, payload: "\x00\x00\x00\x00"))
       }
-      refute_raises {
-        stream.process_frame(Frame.new(type: :headers, stream_id: stream.id))
-      }
     }
   end
 
   def test_stream_handle_data
-    test = -> (state = :open, &blk) {
-      con = open_server_connection
-      stream = Stream.new(con, 2)
-      stream.instance_eval { @state = state }
-      blk.call(stream)
-    }
-
     payload = "ABC" * 5
 
-    test.call {|stream|
+    open_new_stream(:open) {|stream|
       data = nil
       stream.on(:data) {|_data| data = _data }
       stream.process_frame(Frame.new(type: :data, stream_id: stream.id,
@@ -63,7 +36,7 @@ class StreamTest < Minitest::Test
       assert_equal(payload, data)
     }
 
-    test.call {|stream|
+    open_new_stream(:open) {|stream|
       data = nil
       stream.on(:data) {|_data| data = _data }
       stream.process_frame(Frame.new(type: :data, stream_id: stream.id,
@@ -71,37 +44,30 @@ class StreamTest < Minitest::Test
       assert_equal(payload, data)
     }
 
-    test.call {|stream|
+    open_new_stream(:open) {|stream|
       assert_connection_error(:protocol_error) {
         stream.process_frame(Frame.new(type: :data, stream_id: stream.id,
                                        flags: [:padded], payload: "".push_uint8(100).push(payload).push("\x00"*6)))
       }
     }
 
-    test.call {|stream|
+    open_new_stream(:open) {|stream|
       stream.process_frame(Frame.new(type: :data, stream_id: stream.id,
                                      flags: [:end_stream], payload: payload))
       assert_equal(:half_closed_remote, stream.state)
     }
 
-    test.call(:half_closed_remote) {|stream|
+    open_new_stream(:half_closed_remote) {|stream|
       stream.process_frame(Frame.new(type: :data, stream_id: stream.id,
                                      flags: [:end_stream], payload: payload))
-      last = sent_frames(stream.connection).last
+      last = sent_frames.last
       assert_equal(:rst_stream, last.type)
       assert_equal(StreamError.new(:stream_closed).http2_error_code, last.payload.uint32)
     }
   end
 
-  def test_stream_handle_headers
-    test = -> (state = :idle, &blk) {
-      con = open_server_connection
-      stream = Stream.new(con, 2)
-      stream.instance_eval { @state = state }
-      blk.call(stream)
-    }
-
-    test.call {|stream|
+  def test_stream_handle_headers_single
+    open_new_stream {|stream|
       headers = nil
       stream.on(:headers) {|_headers|
         headers = _headers
@@ -113,8 +79,10 @@ class StreamTest < Minitest::Test
       assert_equal(:open, stream.state)
       assert_equal([[":path", "/"]], headers)
     }
+  end
 
-    test.call {|stream|
+  def test_stream_handle_headers_continuation
+    open_new_stream {|stream|
       payload = HPACK::Encoder.new(0).encode([[":path", "/"]])
       headers = nil
       stream.on(:headers) {|_headers|
@@ -132,8 +100,10 @@ class StreamTest < Minitest::Test
       assert_equal(:half_closed_remote, stream.state)
       assert_equal([[":path", "/"]], headers)
     }
+  end
 
-    test.call {|stream|
+  def test_stream_handle_headers_padded
+    open_new_stream {|stream|
       payload = HPACK::Encoder.new(0).encode([[":path", "/"]])
       headers = nil
       stream.on(:headers) {|_headers|
@@ -145,8 +115,10 @@ class StreamTest < Minitest::Test
                                      payload: "".push_uint8(payload.bytesize).push(payload).push("\x00"*payload.bytesize)))
       assert_equal([[":path", "/"]], headers)
     }
+  end
 
-    test.call {|stream|
+  def test_stream_handle_headers_too_long_padding
+    open_new_stream {|stream|
       payload = HPACK::Encoder.new(0).encode([[":path", "/"]])
       assert_connection_error(:protocol_error) {
         stream.process_frame(Frame.new(type: :headers,
@@ -155,8 +127,10 @@ class StreamTest < Minitest::Test
                                        payload: "".push_uint8(payload.bytesize+1).push(payload).push("\x00"*(payload.bytesize+1))))
       }
     }
+  end
 
-    test.call {|stream|
+  def test_stream_handle_headers_broken
+    open_new_stream {|stream|
       payload = "\x00\x01\x02"
       assert_connection_error(:compression_error) {
         stream.process_frame(Frame.new(type: :headers,
@@ -165,61 +139,65 @@ class StreamTest < Minitest::Test
                                        payload: payload))
       }
     }
+  end
 
+  def test_stream_handle_headers_state
     _payload = HPACK::Encoder.new(0).encode([[":path", "/"]])
-    test.call(:reserved_local) {|stream|
+    open_new_stream(:reserved_local) {|stream|
       assert_connection_error(:protocol_error) {
         stream.process_frame(Frame.new(type: :headers, stream_id: stream.id, flags: [:end_headers, :end_stream], payload: _payload))
       }
     }
-    test.call(:closed) {|stream|
+    open_new_stream(:closed) {|stream|
       assert_connection_error(:stream_closed) {
         stream.process_frame(Frame.new(type: :headers, stream_id: stream.id, flags: [:end_headers, :end_stream], payload: _payload))
       }
     }
-    test.call(:half_closed_remote) {|stream|
+    open_new_stream(:half_closed_remote) {|stream|
       stream.process_frame(Frame.new(type: :headers, stream_id: stream.id, flags: [:end_headers, :end_stream], payload: _payload))
-      last = sent_frames(stream.connection).last
+      last = sent_frames.last
       assert_equal(:rst_stream, last.type)
       assert_equal(StreamError.new(:stream_closed).http2_error_code, last.payload.uint32)
     }
+  end
 
-    test.call {|stream|
+  def test_stream_handle_headers_priority
+    open_new_stream {|stream|
       skip "HEADERS with priority" # TODO
     }
   end
 
   def test_stream_promise
-    con = open_server_connection
-    stream = con.__send__(:new_stream, 3)
-    push_stream = stream.promise([])
-    assert(push_stream.id % 2 == 0)
-    assert(push_stream.id > stream.id)
-    assert_equal(stream, push_stream.parent)
-    assert_includes(stream.children, push_stream)
+    open_new_stream {|stream|
+      push_stream = stream.promise([])
+
+      assert(push_stream.id % 2 == 0)
+      assert(push_stream.id > stream.id)
+      assert_equal(stream, push_stream.parent)
+      assert_includes(stream.children, push_stream)
+    }
   end
 
   def test_stream_window_update
-    con = open_server_connection
-    stream = Stream.new(con, 1)
-    before_ws = stream.recv_remaining_window
-    stream.window_update(500)
-
-    last = sent_frames(con).last
-    assert_equal(:window_update, last.type)
-    assert_equal(500, last.payload.uint32)
-    assert_equal(before_ws + 500, stream.recv_remaining_window)
+    open_new_stream {|stream|
+      before_ws = stream.recv_remaining_window
+      stream.window_update(500)
+
+      last = sent_frames.last
+      assert_equal(:window_update, last.type)
+      assert_equal(500, last.payload.uint32)
+      assert_equal(before_ws + 500, stream.recv_remaining_window)
+    }
   end
 
   def test_stream_close
-    con = open_server_connection
-    stream = Stream.new(con, 1)
-    stream.instance_eval { @state = :half_closed_local }
-    stream.close(StreamError.new(:frame_size_error).http2_error_code)
-
-    last = sent_frames(con).last
-    assert_equal(:rst_stream, last.type)
-    assert_equal(StreamError.new(:frame_size_error).http2_error_code, last.payload.uint32)
-    assert_equal(:closed, stream.state)
+    open_new_stream(:half_closed_local) {|stream|
+      stream.close(StreamError.new(:frame_size_error).http2_error_code)
+
+      last = sent_frames.last
+      assert_equal(:rst_stream, last.type)
+      assert_equal(StreamError.new(:frame_size_error).http2_error_code, last.payload.uint32)
+      assert_equal(:closed, stream.state)
+    }
   end
 end
diff --git a/test/utils.rb b/test/utils.rb
index 53aa9e7..cf0e0e5 100644
--- a/test/utils.rb
+++ b/test/utils.rb
@@ -33,18 +33,29 @@ module ServerTestUtils
   private
   def open_server_connection
     io = StringIO.new
-    con = ServerConnection.new(io)
-    con << ServerConnection::CLIENT_CONNECTION_PREFACE
-    con << Frame.new(type: :settings, stream_id: 0).assemble
+    @_con = ServerConnection.new(io)
+    @_con << ServerConnection::CLIENT_CONNECTION_PREFACE
+    @_con << Frame.new(type: :settings, stream_id: 0).assemble
     if block_given?
-      yield con
+      yield @_con
     else
-      con
+      @_con
     end
   end
 
-  def sent_frames(con)
-    resp = con.socket.string.dup
+  def open_new_stream(state = :idle)
+    open_server_connection do |con|
+      @_stream = con.instance_eval { new_stream(3, state: state) }
+      if block_given?
+        yield @_stream
+      else
+        @_stream
+      end
+    end
+  end
+
+  def sent_frames(con = nil)
+    resp = (con || @_con).socket.string.dup
     frames = []
     while f = Frame.parse!(resp)
       frames << f