3 files changed, 38 insertions, 21 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
diff --git a/actionpack/test/controller/live_stream_test.rb b/actionpack/test/controller/live_stream_test.rb
index 934915fc5b..5fb066cc51 100644
--- a/actionpack/test/controller/live_stream_test.rb
+++ b/actionpack/test/controller/live_stream_test.rb
@@ -1,5 +1,6 @@
 require 'abstract_unit'
 require 'active_support/concurrency/latch'
+Thread.abort_on_exception = true
 
 module ActionController
   class SSETest < ActionController::TestCase
@@ -43,9 +44,7 @@ module ActionController
     tests SSETestController
 
     def wait_for_response_stream_close
-      while !response.stream.closed?
-        sleep 0.01
-      end
+      response.stream.await_close
     end
 
     def test_basic_sse
@@ -175,16 +174,6 @@ module ActionController
 
     tests TestController
 
-    class TestResponse < Live::Response
-      def recycle!
-        initialize
-      end
-    end
-
-    def build_response
-      TestResponse.new
-    end
-
     def assert_stream_closed
       assert response.stream.closed?, 'stream should be closed'
     end