Add support for streaming of response body.

This enables testing of infinite Server-Sent Event or RFC-7464 responses.
mdub · Apr 14, 2018 · 4ae9c4b · 4ae9c4b
1 parent 7314c0d
commit 4ae9c4b
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 13 deletions.
diff --git a/CHANGES.md b/CHANGES.md
@@ -1,3 +1,7 @@
+## PENDING (14-Apr-2018)
+
+* Add support for streaming of response body with `Net:HTTP` and friends.
+
 ## 1.4.1 (20-Aug-2017)
 
 * Add support for `Net::HTTPRequest#body_stream`, and with it, newer versions of rest-client.

diff --git a/lib/sham_rack/net_http.rb b/lib/sham_rack/net_http.rb
@@ -81,31 +81,25 @@ def build_response(rack_response)
         code, message = status.to_s.split(" ", 2)
         message ||= Rack::Utils::HTTP_STATUS_CODES[code.to_i]
         response = Net::HTTPResponse.send(:response_class, code).new("Sham", code, message)
-        response.instance_variable_set(:@body, assemble_body(body))
-        response.instance_variable_set(:@read, true)
+        response.instance_variable_set(:@rack_body, body)
         headers.each do |k,v|
           response.add_field(k, v)
         end
         response.extend ShamRack::NetHttp::ResponseExtensions
         return response
       end
 
-      def assemble_body(body)
-        content = ""
-        body.each { |fragment| content << fragment }
-        content
-      end
-
     end
 
     module ResponseExtensions
 
-      def read_body(dest = nil)
-        yield @body if block_given?
-        dest << @body if dest
-        return @body
+      def read_body(dest = nil, &block)
+        out = procdest(dest, block)
+        @rack_body.each do |fragment|
+          out << fragment
+        end
+        out
       end
-
     end
 
   end

diff --git a/spec/sham_rack_spec.rb b/spec/sham_rack_spec.rb
@@ -233,6 +233,49 @@ class NetHttpProhibited < StandardError; end
 
   end
 
+  describe "streaming response" do
+
+    class Countdown
+      def each
+        10.downto(1) do |s|
+          yield "#{s}\n"
+        end
+        raise "BOOM!"
+      end
+    end
+
+    before(:each) do
+      ShamRack.at("www.bombs-r-us.com") do
+        [
+          "200",
+          { "Content-Type" => "text/plain" },
+          Countdown.new
+        ]
+      end
+    end
+
+    it "can be streamed using Net::HTTP" do
+      response = Net::HTTP.start("www.bombs-r-us.com") do |http|
+        http.request(Net::HTTP::Get.new("/")) do |response|
+          chunks = response.enum_for(:read_body)
+          expect(chunks.take(3).to_a).to eq(["10\n", "9\n", "8\n"])
+        end
+      end
+    end
+
+    it "can be streamed using RestClient" do
+      chunks = []
+      handler = proc do |response|
+        response.enum_for(:read_body).take(3).each do |chunk|
+          chunks << chunk
+        end
+      end
+      RestClient::Request.execute(:method => :get, :url => 'http://www.bombs-r-us.com', :block_response => handler)
+      expect(chunks).to eq(["10\n", "9\n", "8\n"])
+    end
+
+  end
+
   describe ".allow_network_connections" do
 
     context "when false" do