update tests

4 files changed, 45 insertions, 23 deletions

diff --git a/lib/plum/client.rb b/lib/plum/client.rb
index fa5bdd8..3749482 100644
--- a/lib/plum/client.rb
+++ b/lib/plum/client.rb
@@ -48,19 +48,23 @@ module Plum
     # @param response [Response] if specified, waits only for the response
     def wait(response = nil)
       if response
-        _succ while !response.finished?
+        _succ while !response.failed? && !response.finished?
       else
         _succ while !@responses.empty?
       end
     end
 
+    # Waits for the response headers.
+    # @param response [Response] the incomplete response.
+    def wait_headers(response)
+      _succ while !response.failed? && !response.headers
+    end
+
     # Closes the connection.
     def close
-      begin
-        @plum.close if @plum
-      ensure
-        @socket.close if @socket
-      end
+      @plum.close if @plum && @plum.state != :closed
+    ensure
+      @socket.close if @socket
     end
 
     # Creates a new HTTP request.
@@ -169,18 +173,22 @@ module Plum
       }
       plum = ClientConnection.new(sock.method(:write), local_settings)
       plum.on(:protocol_error) { |ex|
-        _fail
+        _fail(ex)
         raise ex
       }
+      plum.on(:close) { _fail(RuntimeError.new(:closed)) }
       plum.on(:stream_error) { |stream, ex|
         if res = @responses.delete(stream)
-          res._fail unless res.finished?
+          res._fail(ex) unless res.finished?
         end
         raise ex
       }
       plum.on(:headers) { |stream, headers|
         response = @responses[stream]
         response._headers(headers)
+        if handler = @response_handlers.delete(stream)
+          handler.call(response)
+        end
       }
       plum.on(:data) { |stream, chunk|
         response = @responses[stream]
@@ -189,21 +197,18 @@ module Plum
       plum.on(:end_stream) { |stream|
         response = @responses.delete(stream)
         response._finish
-        if handler = @response_handlers.delete(stream)
-          handler.call(response)
-        end
       }
       plum
     end
 
     def _succ
-      _fail if @socket.closed? || @socket.eof?
       @plum << @socket.readpartial(1024)
     end
 
-    def _fail
-      while res = @responses.pop
-        res._fail unless res.finished?
+    def _fail(ex)
+      while sr = @responses.shift
+        stream, res = sr
+        res._fail(ex) unless res.finished?
       end
     ensure
       close
diff --git a/lib/plum/client/response.rb b/lib/plum/client/response.rb
index 6afa4f9..ee23971 100644
--- a/lib/plum/client/response.rb
+++ b/lib/plum/client/response.rb
@@ -43,9 +43,9 @@ module Plum
     def each_chunk(&block)
       raise "Body already read" if @body_read
       @body_read = true
-      while chunk = @body.pop
-        if chunk == :failed
-          raise EOFError
+      while !(finished? && @body.empty?) && chunk = @body.pop
+        if Exception === chunk
+          raise chunk
         else
           yield chunk
         end
@@ -78,9 +78,9 @@ module Plum
     end
 
     # @api private
-    def _fail
+    def _fail(ex)
       @failed = true
-      @body << :failed
+      @body << ex
     end
   end
 end
diff --git a/lib/plum/connection.rb b/lib/plum/connection.rb
index 68ba07d..f365861 100644
--- a/lib/plum/connection.rb
+++ b/lib/plum/connection.rb
@@ -38,6 +38,7 @@ module Plum
 
     # Emits :close event. Doesn't actually close socket.
     def close
+      @state = :closed
       # TODO: server MAY wait streams
       callback(:close)
     end
@@ -45,6 +46,7 @@ module Plum
     # Receives the specified data and process.
     # @param new_data [String] The data received from the peer.
     def receive(new_data)
+      return if @state == :closed
       return if new_data.empty?
       @buffer << new_data
       consume_buffer
@@ -135,9 +137,7 @@ module Plum
       when :ping
         receive_ping(frame)
       when :goaway
-        callback(:goaway, frame)
-        goaway
-        close
+        receive_goaway(frame)
       when :data, :headers, :priority, :rst_stream, :push_promise, :continuation
         raise Plum::ConnectionError.new(:protocol_error)
       else
@@ -184,5 +184,18 @@ module Plum
         send_immediately Frame.ping(:ack, opaque_data)
       end
     end
+
+    def receive_goaway(frame)
+      callback(:goaway, frame)
+      goaway
+      close
+
+      last_id = frame.payload.uint32(0)
+      error_code = frame.payload.uint32(4)
+      message = frame.payload.byteslice(8, frame.length - 8)
+      if error_code > 0
+        raise LocalConnectionError.new(HTTPError::ERROR_CODES.key(error_code), message)
+      end
+    end
   end
 end
diff --git a/lib/plum/errors.rb b/lib/plum/errors.rb
index 9df4668..d5cce48 100644
--- a/lib/plum/errors.rb
+++ b/lib/plum/errors.rb
@@ -43,4 +43,8 @@ module Plum
       @parser = parser
     end
   end
+
+  # Client
+  class LocalConnectionError < HTTPError; end
+  class LocalStreamError < HTTPError; end
 end