diff options
-rw-r--r-- | lib/plum/client/response.rb | 12 | ||||
-rw-r--r-- | test/plum/client/test_response.rb | 11 |
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 |