2 files changed, 36 insertions, 8 deletions
diff --git a/actionpack/lib/action_controller/metal/live.rb b/actionpack/lib/action_controller/metal/live.rb
index 5ef4f6ccda..fba17746c0 100644
--- a/actionpack/lib/action_controller/metal/live.rb
+++ b/actionpack/lib/action_controller/metal/live.rb
@@ -107,8 +107,11 @@ module ActionController
     end
 
     class Buffer < ActionDispatch::Response::Buffer #:nodoc:
+      include MonitorMixin
+
       def initialize(response)
         @error_callback = lambda { true }
+        @cv = new_cond
         super(response, SizedQueue.new(10))
       end
 
@@ -128,8 +131,17 @@ module ActionController
       end
 
       def close
-        super
-        @buf.push nil
+        synchronize do
+          super
+          @buf.push nil
+          @cv.broadcast
+        end
+      end
+
+      def await_close
+        synchronize do
+          @cv.wait_until { @closed }
+        end
       end
 
       def on_error(&block)
diff --git a/actionpack/lib/action_controller/test_case.rb b/actionpack/lib/action_controller/test_case.rb
index 8650b75400..33a5858766 100644
--- a/actionpack/lib/action_controller/test_case.rb
+++ b/actionpack/lib/action_controller/test_case.rb
@@ -258,6 +258,17 @@ module ActionController
     end
   end
 
+  class LiveTestResponse < Live::Response
+    def recycle!
+      @body = nil
+      initialize
+    end
+
+    def body
+      @body ||= super
+    end
+  end
+
   # Methods #destroy and #load! are overridden to avoid calling methods on the
   # @store object, which does not exist for the TestSession class.
   class TestSession < Rack::Session::Abstract::SessionHash #:nodoc:
@@ -583,13 +594,14 @@ module ActionController
       end
 
       def setup_controller_request_and_response
-        @request          = build_request
-        @response         = build_response
-        @response.request = @request
-
         @controller = nil unless defined? @controller
 
+        response_klass = TestResponse
+
         if klass = self.class.controller_class
+          if klass < ActionController::Live
+            response_klass = LiveTestResponse
+          end
           unless @controller
             begin
               @controller = klass.new
@@ -599,6 +611,10 @@ module ActionController
           end
         end
 
+        @request          = build_request
+        @response         = build_response response_klass
+        @response.request = @request
+
         if @controller
           @controller.request = @request
           @controller.params = {}
@@ -609,8 +625,8 @@ module ActionController
         TestRequest.new
       end
 
-      def build_response
-        TestResponse.new
+      def build_response(klass)
+        klass.new
       end
 
       included do