client/response: add Response#on_finish

1 files changed, 12 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