client/response: add Response#on_finish

2 files changed, 23 insertions, 0 deletions

diff --git a/lib/plum/client/response.rb b/lib/plum/client/response.rb
index 5c0b2a0..e9bd5e1 100644
--- a/lib/plum/client/response.rb
+++ b/lib/plum/client/response.rb
@@ -42,6 +42,7 @@ module Plum
     # @yield [chunk] A chunk of the response body.
     def on_chunk(&block)
       raise "Body already read" if @on_chunk
+      raise ArgumentError, "block must be given" unless block_given?
       @on_chunk = block
       unless @body.empty?
         @body.each(&block)
@@ -49,6 +50,16 @@ module Plum
       end
     end
 
+    # Set callback that will be called when the response finished.
+    def on_finish(&block)
+      raise ArgumentError, "block must be given" unless block_given?
+      if finished?
+        block.call
+      else
+        @on_finish = block
+      end
+    end
+
     # Returns the complete response body. Use #each_body instead if the body can be very large.
     # @return [String] the whole response body
     def body
@@ -78,6 +89,7 @@ module Plum
     # @api private
     def _finish
       @finished = true
+      @on_finish.call if @on_finish
     end
 
     # @api private
diff --git a/test/plum/client/test_response.rb b/test/plum/client/test_response.rb
index 553bf45..76d9037 100644
--- a/test/plum/client/test_response.rb
+++ b/test/plum/client/test_response.rb
@@ -60,4 +60,15 @@ class ResponseTest < Minitest::Test
     resp._chunk("c")
     assert_equal(["a", "b", "c"], res)
   end
+
+  def test_on_finish
+    resp = Response.new
+    ran = false
+    resp.on_finish { ran = true }
+    resp._finish
+    assert(ran)
+    ran = false
+    resp.on_finish { ran = true }
+    assert(ran)
+  end
 end