live controllers should have live responses

detect the type of controller we're testing and return the right type of
response based on that controller.  This allows us to stop doing the
weird sleep thing.

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