1 files changed, 68 insertions, 17 deletions

diff --git a/actionpack/test/controller/live_stream_test.rb b/actionpack/test/controller/live_stream_test.rb
index 0a431270b5..947f64176b 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.body
     end
 
     def test_basic_sse
@@ -91,6 +90,9 @@ module ActionController
   end
 
   class LiveStreamTest < ActionController::TestCase
+    class Exception < StandardError
+    end
+
     class TestController < ActionController::Base
       include ActionController::Live
 
@@ -100,6 +102,12 @@ module ActionController
         'test'
       end
 
+      def set_cookie
+        cookies[:hello] = "world"
+        response.stream.write "hello world"
+        response.close
+      end
+
       def render_text
         render :text => 'zomg'
       end
@@ -145,6 +153,11 @@ module ActionController
         render 'doesntexist'
       end
 
+      def exception_in_view_after_commit
+        response.stream.write ""
+        render 'doesntexist'
+      end
+
       def exception_with_callback
         response.headers['Content-Type'] = 'text/event-stream'
 
@@ -153,32 +166,33 @@ module ActionController
           response.stream.close
         end
 
+        response.stream.write "" # make sure the response is committed
         raise 'An exception occurred...'
       end
 
+      def exception_in_controller
+        raise Exception, 'Exception in controller'
+      end
+
+      def bad_request_error
+        raise ActionController::BadRequest
+      end
+
       def exception_in_exception_callback
         response.headers['Content-Type'] = 'text/event-stream'
         response.stream.on_error do
           raise 'We need to go deeper.'
         end
+        response.stream.write ''
         response.stream.write params[:widget][:didnt_check_for_nil]
       end
     end
 
     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'
+      assert response.sent?, 'stream should be sent'
     end
 
     def capture_log_output
@@ -192,6 +206,13 @@ module ActionController
       end
     end
 
+    def test_set_cookie
+      @controller = TestController.new
+      get :set_cookie
+      assert_equal({'hello' => 'world'}, @response.cookies)
+      assert_equal "hello world", @response.body
+    end
+
     def test_set_response!
       @controller.set_response!(@request)
       assert_kind_of(Live::Response, @controller.response)
@@ -213,6 +234,7 @@ module ActionController
       @controller.response = @response
 
       t = Thread.new(@response) { |resp|
+        resp.await_commit
         resp.stream.each do |part|
           assert_equal parts.shift, part
           ol = @controller.latch
@@ -249,24 +271,34 @@ module ActionController
     end
 
     def test_exception_handling_html
-      capture_log_output do |output|
+      assert_raises(ActionView::MissingTemplate) do
         get :exception_in_view
+      end
+
+      capture_log_output do |output|
+        get :exception_in_view_after_commit
         assert_match %r((window\.location = "/500\.html"</script></html>)$), response.body
         assert_match 'Missing template test/doesntexist', output.rewind && output.read
         assert_stream_closed
       end
+      assert response.body
+      assert_stream_closed
     end
 
     def test_exception_handling_plain_text
-      capture_log_output do |output|
+      assert_raises(ActionView::MissingTemplate) do
         get :exception_in_view, format: :json
+      end
+
+      capture_log_output do |output|
+        get :exception_in_view_after_commit, format: :json
         assert_equal '', response.body
         assert_match 'Missing template test/doesntexist', output.rewind && output.read
         assert_stream_closed
       end
     end
 
-    def test_exception_callback
+    def test_exception_callback_when_committed
       capture_log_output do |output|
         get :exception_with_callback, format: 'text/event-stream'
         assert_equal %(data: "500 Internal Server Error"\n\n), response.body
@@ -275,7 +307,19 @@ module ActionController
       end
     end
 
-    def test_exceptions_raised_handling_exceptions
+    def test_exception_in_controller_before_streaming
+      assert_raises(ActionController::LiveStreamTest::Exception) do
+        get :exception_in_controller, format: 'text/event-stream'
+      end
+    end
+
+    def test_bad_request_in_controller_before_streaming
+      assert_raises(ActionController::BadRequest) do
+        get :bad_request_error, format: 'text/event-stream'
+      end
+    end
+
+    def test_exceptions_raised_handling_exceptions_and_committed
       capture_log_output do |output|
         get :exception_in_exception_callback, format: 'text/event-stream'
         assert_equal '', response.body
@@ -295,4 +339,11 @@ module ActionController
       assert_equal 304, @response.status.to_i
     end
   end
+
+  class BufferTest < ActionController::TestCase
+    def test_nil_callback
+      buf = ActionController::Live::Buffer.new nil
+      assert buf.call_on_error
+    end
+  end
 end