Preserve one-shot response bodies after validation

lib/committee/middleware/response_validation.rb

@@ -28,8 +28,19 @@ def handle(request)
             end
           end
         else
+          validator = response_validator(request, status)
+          if validator
+            original_response = response
+            response = []
+            begin
+              original_response.each { |chunk| response << chunk }
+            ensure
+              original_response.close if original_response.respond_to?(:close)
+            end
+          end
+
           begin
-            validate(request, status, headers, response)
+            validator&.response_validate(status, headers, response, @strict)
           rescue Committee::InvalidResponse
             handle_exception($!, request.env)
 
@@ -68,10 +79,12 @@ def build_options(options)
       end
 
       def validate(request, status, headers, response, streaming_content_parser = nil)
-        v = build_schema_validator(request)
-        if v.link_exist? && self.class.validate?(status, validate_success_only)
-          v.response_validate(status, headers, response, @strict, streaming_content_parser)
-        end
+        response_validator(request, status)&.response_validate(status, headers, response, @strict, streaming_content_parser)
+      end
+
+      def response_validator(request, status)
+        validator = build_schema_validator(request)
+        validator if validator.link_exist? && self.class.validate?(status, validate_success_only)
       end
 
       def retrieve_streaming_content_parser(headers)

test/middleware/response_validation_open_api_3_test.rb

@@ -17,6 +17,29 @@ def app
     assert_equal 200, last_response.status
   end
 
+  it "preserves a one-shot response body after validation" do
+    content = JSON.generate(CHARACTERS_RESPONSE)
+    chunks = [content]
+    each_calls = 0
+    body = Rack::BodyProxy.new(Enumerator.new do |yielder|
+      each_calls += 1
+      chunks.each { |chunk| yielder << chunk }
+      chunks.clear
+    end) {}
+    @app = Rack::Builder.new {
+      use Committee::Middleware::ResponseValidation, { schema: open_api_3_schema }
+      run ->(_) { [200, { "Content-Type" => "application/json" }, body] }
+    }
+
+    status, _headers, response_body = @app.call(Rack::MockRequest.env_for("/characters"))
+
+    assert_equal 200, status
+    assert_equal content, response_body.each.to_a.join
+    assert_equal 1, each_calls
+    response_body.close if response_body.respond_to?(:close)
+    assert body.closed?
+  end
+
   it "passes through a valid response with content-type (lower-case)" do
     status = 200
     headers = { "content-type" => "application/json" }