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.

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