1 files changed, 27 insertions, 28 deletions

diff --git a/actionpack/test/controller/live_stream_test.rb b/actionpack/test/controller/live_stream_test.rb
index fb6a750089..eb8b2c9832 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
@@ -91,6 +90,9 @@ module ActionController
   end
 
   class LiveStreamTest < ActionController::TestCase
+    class Exception < StandardError
+    end
+
     class TestController < ActionController::Base
       include ActionController::Live
 
@@ -153,11 +155,12 @@ 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 in controller'
+        raise Exception, 'Exception in controller'
       end
 
       def bad_request_error
@@ -169,23 +172,15 @@ module ActionController
         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
+      response.stream.await_close
       assert response.stream.closed?, 'stream should be closed'
     end
 
@@ -257,24 +252,19 @@ module ActionController
     end
 
     def test_exception_handling_html
-      capture_log_output do |output|
+      assert_raises(ActionView::MissingTemplate) do
         get :exception_in_view
-        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_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
-        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
@@ -284,16 +274,18 @@ module ActionController
     end
 
     def test_exception_in_controller_before_streaming
-      response = get :exception_in_controller, format: 'text/event-stream'
-      assert_equal 500, response.status
+      assert_raises(ActionController::LiveStreamTest::Exception) do
+        get :exception_in_controller, format: 'text/event-stream'
+      end
     end
 
     def test_bad_request_in_controller_before_streaming
-      response = get :bad_request_error, format: 'text/event-stream'
-      assert_equal 400, response.status
+      assert_raises(ActionController::BadRequest) do
+        get :bad_request_error, format: 'text/event-stream'
+      end
     end
 
-    def test_exceptions_raised_handling_exceptions
+    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
@@ -313,4 +305,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