restructure core extensions and stub to fix issues with live requests

spec/webmock_spec.cr

@@ -470,7 +470,6 @@ describe WebMock 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
 

src/webmock/core_ext.cr

@@ -7,37 +7,69 @@ 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
+    before_request(request)
+    WebMock.find_stub(request).try { |stub| return stub.exec(request) }
+
+    perform_live_request(request)
   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)
+  private def exec_internal(request : HTTP::Request, &block : Response -> T) forall T
+    before_request(request)
+    WebMock.find_stub(request).try do |stub|
+      stub.as_block
+      return yield(stub.exec(request))
+    end
 
-    stub = WebMock.find_stub(request)
-    return yield(stub.exec(request)) if stub
-    raise WebMock::NetConnectNotAllowedError.new(request) unless WebMock.allows_net_connect?
+    perform_live_request(request) do |response|
+      yield response
+    end
+  end
 
-    request.headers["User-agent"] ||= "Crystal"
-    request.to_io(socket)
-    socket.flush
+  private def perform_live_request(request)
+    send_live_request(request)
+    receive_live_response(request)
+  end
+
+  private def perform_live_request(request, &block)
+    send_live_request(request)
+    receive_live_response(request) do |response|
+      yield response
+    end
+  end
 
-    result = nil
+  private def receive_live_response(request)
+    HTTP::Client::Response.from_io(socket, request.ignore_body?).tap do |response|
+      after_live_request(request, response)
+    end
+  end
 
+  private def receive_live_response(request, &block)
     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)
+      yield(response).tap do
+        after_live_request(request, response)
+      end
     end
+  end
+
+  private def after_live_request(request, 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)
+  private def send_live_request(request)
+    unless WebMock.allows_net_connect?
+      raise WebMock::NetConnectNotAllowedError.new(request)
+    end
 
-    result
+    request.headers["User-agent"] ||= "Crystal"
+    request.to_io(socket)
+    socket.flush
+  end
+
+  private def before_request(request)
+    request.scheme = "https" if tls?
+    request.headers["Host"] = host_header unless request.headers.has_key?("Host")
+    run_before_request_callbacks(request)
   end
 end

src/webmock/stub.cr

@@ -3,7 +3,9 @@ class WebMock::Stub
   @uri : URI
   @expected_headers : HTTP::Headers?
   @calls = 0
+  @body : String
   @body_io : IO?
+  @requested_as_block = false
 
   def initialize(method : Symbol | String, uri)
     @method = method.to_s.upcase
@@ -15,6 +17,8 @@ class WebMock::Stub
     @headers = HTTP::Headers{"Content-length" => "0", "Connection" => "close"}
 
     @block = Proc(HTTP::Request, HTTP::Client::Response).new do |request|
+      copy_body unless @requested_as_block
+      set_content_headers
       HTTP::Client::Response.new(@status, body: @body, headers: @headers, body_io: @body_io)
     end
   end
@@ -30,24 +34,25 @@ class WebMock::Stub
     @body = body
     @body_io = nil
     @status = status
-    @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 = ""
     @body_io = body_io
     @status = status
-    @headers.delete("Content-length")
-    @headers["Transfer-encoding"] = "chunked"
     @headers.merge!(headers) if headers
     self
   end
 
   def to_return(&block : HTTP::Request -> HTTP::Client::Response)
-    @block = block
+    @block = Proc(HTTP::Request, HTTP::Client::Response).new do |request|
+      copy_body unless @requested_as_block
+      set_content_headers
+      block.call(request)
+    end
+
     self
   end
 
@@ -121,6 +126,24 @@ class WebMock::Stub
     @calls
   end
 
+  def as_block
+    @requested_as_block = true
+  end
+
+  private def set_content_headers
+    if @requested_as_block
+      @headers.delete("Content-Length")
+      @headers["Transfer-Encoding"] = "chunked"
+    else
+      @headers["Content-Length"] = @body.bytesize.to_s
+      @headers.delete("Transfer-Encoding")
+    end
+  end
+
+  private def copy_body
+    @body_io.try { |io| @body = io.gets_to_end || "" }
+  end
+
   private def parse_uri(uri_string)
     uri = URI.parse(uri_string)
     uri = URI.parse("http://#{uri_string}") unless uri.host