From 78bb0e3b5850c700da0e7fbdd2d6c180cc4a061b Mon Sep 17 00:00:00 2001
From: Joakim Repomaa <mail@jreinert.com>
Date: Mon, 7 Oct 2019 19:19:13 +0200
Subject: [PATCH] Add support for yielding HTTP::Client method types (#27)

---
 spec/webmock_spec.cr    | 57 +++++++++++++++++++++++++++++++++++++++++
 src/webmock/core_ext.cr | 39 ++++++++++++++++++----------
 src/webmock/stub.cr     | 19 +++++++++++---
 3 files changed, 99 insertions(+), 16 deletions(-)

diff --git a/spec/webmock_spec.cr b/spec/webmock_spec.cr
index e59d6ae..12f4182 100644
--- a/spec/webmock_spec.cr
+++ b/spec/webmock_spec.cr
@@ -464,6 +464,63 @@ describe WebMock do
     end
   end
 
+  context "with body_io" do
+    it "allows reading the body" do
+      WebMock.wrap do
+        WebMock.stub(:get, "http://www.example.com").to_return(body_io: IO::Memory.new("Hello!"))
+
+        body = HTTP::Client.get("http://www.example.com").body
+        body.should eq("Hello!")
+      end
+    end
+
+    it "sets content-length header correctly" do
+      WebMock.wrap do
+        WebMock.stub(:get, "http://www.example.com").to_return(body_io: IO::Memory.new("Hello!"))
+
+        headers = HTTP::Client.get("http://www.example.com").headers
+        headers["Content-length"].should eq("6")
+      end
+    end
+  end
+
+  context "with yielding method variants" do
+    it "stubs the request" do
+      WebMock.wrap do
+        stub = WebMock.stub(:get, "http://www.example.com")
+
+        HTTP::Client.get("http://www.example.com") do |response|
+        end
+
+        stub.calls.should eq(1)
+      end
+    end
+
+    it "allows setting and reading body_io" do
+      WebMock.wrap do
+        WebMock.stub(:get, "http://www.example.com").to_return(body_io: IO::Memory.new("Hello!"))
+
+        body = HTTP::Client.get("http://www.example.com") do |response|
+          response.body_io.gets_to_end
+        end
+
+        body.should eq("Hello!")
+      end
+    end
+
+    it "sets transfer encoding header to chunked" do
+      WebMock.wrap do
+        WebMock.stub(:get, "http://www.example.com").to_return(body_io: IO::Memory.new("Hello!"))
+
+        headers = HTTP::Client.get("http://www.example.com") do |response|
+          response.headers
+        end
+
+        headers.not_nil!["Transfer-encoding"].should eq("chunked")
+      end
+    end
+  end
+
   # Commented so that specs run fast, but uncomment to try it (it works)
   # it "calls callback after live request" do
   #  WebMock.wrap do
diff --git a/src/webmock/core_ext.cr b/src/webmock/core_ext.cr
index 0a59655..c0e6e4a 100644
--- a/src/webmock/core_ext.cr
+++ b/src/webmock/core_ext.cr
@@ -7,24 +7,37 @@ end
 
 class HTTP::Client
   private def exec_internal(request : HTTP::Request)
+    response = exec_internal(request) { |res| res }
+    response.tap do |response|
+      response.consume_body_io
+      response.headers.delete("Transfer-encoding")
+      response.headers["Content-length"] = response.body.bytesize.to_s
+    end
+  end
+
+  private def exec_internal(request : HTTP::Request, &block : Response -> T) : T forall T
     request.scheme = "https" if tls?
     request.headers["Host"] = host_header unless request.headers.has_key?("Host")
     run_before_request_callbacks(request)
 
     stub = WebMock.find_stub(request)
-    return stub.exec(request) if stub
-
-    if WebMock.allows_net_connect?
-      request.headers["User-agent"] ||= "Crystal"
-      request.to_io(socket)
-      socket.flush
-      res = HTTP::Client::Response.from_io(socket, request.ignore_body?).tap do |response|
-        close unless response.keep_alive?
-      end
-      WebMock.callbacks.call(:after_live_request, request, res)
-      res
-    else
-      raise WebMock::NetConnectNotAllowedError.new(request)
+    return yield(stub.exec(request)) if stub
+    raise WebMock::NetConnectNotAllowedError.new(request) unless WebMock.allows_net_connect?
+
+    request.headers["User-agent"] ||= "Crystal"
+    request.to_io(socket)
+    socket.flush
+
+    result = nil
+
+    HTTP::Client::Response.from_io(socket, request.ignore_body?) do |response|
+      result = yield(response)
+      close unless response.keep_alive?
+      WebMock.callbacks.call(:after_live_request, request, response)
     end
+
+    raise "Unexpected end of response" unless result.is_a?(T)
+
+    result
   end
 end
diff --git a/src/webmock/stub.cr b/src/webmock/stub.cr
index 6e9af27..63c9444 100644
--- a/src/webmock/stub.cr
+++ b/src/webmock/stub.cr
@@ -3,6 +3,7 @@ class WebMock::Stub
   @uri : URI
   @expected_headers : HTTP::Headers?
   @calls = 0
+  @body_io : IO?
 
   def initialize(method : Symbol | String, uri)
     @method = method.to_s.upcase
@@ -14,7 +15,7 @@ class WebMock::Stub
     @headers = HTTP::Headers{"Content-length" => "0", "Connection" => "close"}
 
     @block = Proc(HTTP::Request, HTTP::Client::Response).new do |request|
-      HTTP::Client::Response.new(@status, body: @body, headers: @headers)
+      HTTP::Client::Response.new(@status, body: @body, headers: @headers, body_io: @body_io)
     end
   end
 
@@ -25,10 +26,22 @@ class WebMock::Stub
     self
   end
 
-  def to_return(body = "", status = 200, headers = nil)
+  def to_return(body : String? = "", status = 200, headers = nil)
     @body = body
+    @body_io = nil
     @status = status
-    @headers["Content-length"] = @body.size.to_s
+    @headers.delete("Transfer-encoding")
+    @headers["Content-length"] = body.size.to_s
+    @headers.merge!(headers) if headers
+    self
+  end
+
+  def to_return(body_io : IO, status = 200, headers = nil)
+    @body = nil
+    @body_io = body_io
+    @status = status
+    @headers.delete("Content-length")
+    @headers["Transfer-encoding"] = "chunked"
     @headers.merge!(headers) if headers
     self
   end