client: call the block passed to Client#request in Response#set_headers

This reduces code lines.

4 files changed, 35 insertions, 12 deletions

diff --git a/lib/plum/client/client_session.rb b/lib/plum/client/client_session.rb
index bfc24a7..0c7e366 100644
--- a/lib/plum/client/client_session.rb
+++ b/lib/plum/client/client_session.rb
@@ -42,7 +42,7 @@ module Plum
                   ":scheme" => @config[:https] ? "https" : "http",
       }.merge(headers)
 
-      response = Response.new(**options)
+      response = Response.new(**options, &headers_cb)
       @responses << response
       stream = @plum.open_stream
       stream.send_headers(headers, end_stream: !body)
@@ -50,7 +50,6 @@ module Plum
 
       stream.on(:headers) { |resp_headers_raw|
         response.send(:set_headers, resp_headers_raw.to_h)
-        headers_cb.call(response) if headers_cb
       }
       stream.on(:data) { |chunk|
         response.send(:add_chunk, chunk)
diff --git a/lib/plum/client/legacy_client_session.rb b/lib/plum/client/legacy_client_session.rb
index 166c0ef..07ae6a9 100644
--- a/lib/plum/client/legacy_client_session.rb
+++ b/lib/plum/client/legacy_client_session.rb
@@ -11,7 +11,6 @@ module Plum
       @parser = setup_parser
       @requests = []
       @response = nil
-      @headers_callback = nil
     end
 
     def succ
@@ -40,8 +39,8 @@ module Plum
         end
       end
 
-      response = Response.new(**options)
-      @requests << [response, headers, body, chunked, headers_cb]
+      response = Response.new(**options, &headers_cb)
+      @requests << [response, headers, body, chunked]
       consume_queue
       response
     end
@@ -55,9 +54,8 @@ module Plum
     def consume_queue
       return if @response || @requests.empty?
 
-      response, headers, body, chunked, cb = @requests.shift
+      response, headers, body, chunked = @requests.shift
       @response = response
-      @headers_callback = cb
 
       @socket << construct_request(headers)
 
@@ -98,7 +96,6 @@ module Plum
         # FIXME: duplicate header name?
         resp_headers = parser.headers.map { |key, value| [key.downcase, value] }.to_h
         @response.send(:set_headers, { ":status" => parser.status_code.to_s }.merge(resp_headers))
-        @headers_callback.call(@response) if @headers_callback
       }
 
       parser.on_body = proc { |chunk|
@@ -108,7 +105,6 @@ module Plum
       parser.on_message_complete = proc { |env|
         @response.send(:finish)
         @response = nil
-        @headers_callback = nil
         close unless parser.keep_alive?
         consume_queue
       }
diff --git a/lib/plum/client/response.rb b/lib/plum/client/response.rb
index 0e83fd7..a091c7a 100644
--- a/lib/plum/client/response.rb
+++ b/lib/plum/client/response.rb
@@ -6,19 +6,21 @@ module Plum
     attr_reader :headers
 
     # @api private
-    def initialize(auto_decode: true, **options)
+    def initialize(auto_decode: true, **options, &on_headers)
+      @headers = nil
       @body = Queue.new
       @finished = false
       @failed = false
       @body = []
       @auto_decode = auto_decode
+      @on_headers = on_headers
       @on_chunk = @on_finish = nil
     end
 
     # Returns the HTTP status code.
     # @return [String] the HTTP status code
     def status
-      @headers && @headers[":status"]
+      @headers&.fetch(":status")
     end
 
     # Returns the header value that correspond to the header name.
@@ -40,7 +42,15 @@ module Plum
       @failed
     end
 
-    # Set callback tha called when received a chunk of response body.
+    # Set callback that will be called when the response headers arrive
+    # @yield [self]
+    def on_headers(&block)
+      raise ArgumentError, "block must be given" unless block_given?
+      @on_headers = block
+      yield self if @headers
+    end
+
+    # Set callback that will be called when received a chunk of response body.
     # @yield [chunk] A chunk of the response body.
     def on_chunk(&block)
       raise "Body already read" if @on_chunk
@@ -74,6 +84,7 @@ module Plum
     # internal: set headers and setup decoder
     def set_headers(headers)
       @headers = headers.freeze
+      @on_headers.call(self) if @on_headers
       @decoder = setup_decoder
     end
 
diff --git a/test/plum/client/test_response.rb b/test/plum/client/test_response.rb
index d3d4e42..a1af31c 100644
--- a/test/plum/client/test_response.rb
+++ b/test/plum/client/test_response.rb
@@ -52,6 +52,23 @@ class ResponseTest < Minitest::Test
     }
   end
 
+  def test_on_headers_initialize
+    called = false
+    resp = Response.new { |r| called = true }
+    assert(!called)
+    resp.send(:set_headers, { ":status" => 201 })
+    assert(called)
+  end
+
+  def test_on_headers_explicit
+    called = false
+    resp = Response.new
+    resp.on_headers { |r| called = true }
+    assert(!called)
+    resp.send(:set_headers, { ":status" => 201 })
+    assert(called)
+  end
+
   def test_on_chunk
     resp = Response.new
     resp.send(:set_headers, {})