Fixes #23964

  - Adds #each_chunk to ActionDispatch::Response. it's a method which
    will be called by ActionDispatch::Response#each.
  - Make Response#each a proper method instead of delegating to @stream
  - In Live, instead of overriding #each, override #each_chunk.
  - `#each` should just spit out @str_body if it's already set
  - Adds #test_set_header_after_read_body_during_action
    to prove this fixes #23964
  - Adds #test_each_isnt_called_if_str_body_is_written to
    ensure #each_chunk is not called when @str_body is available
  - Call `@response.sent!` in AC::TestCase's #perform so a test response
    acts a bit more like a real response. Makes test that call  `#assert_stream_closed`
    pass again.
  - Additionally assert `#committed?` in `#assert_stream_closed`
  - Make test that was calling @response.stream.each pass again by
    calling @response.each instead.

1 files changed, 24 insertions, 9 deletions

diff --git a/actionpack/lib/action_dispatch/http/response.rb b/actionpack/lib/action_dispatch/http/response.rb
index fa4c54701a..1515d59df3 100644
--- a/actionpack/lib/action_dispatch/http/response.rb
+++ b/actionpack/lib/action_dispatch/http/response.rb
@@ -68,7 +68,13 @@ module ActionDispatch # :nodoc:
     alias_method :headers,  :header
 
     delegate :[], :[]=, :to => :@header
-    delegate :each, :to => :@stream
+
+    def each(&block)
+      sending!
+      x = @stream.each(&block)
+      sent!
+      x
+    end
 
     CONTENT_TYPE = "Content-Type".freeze
     SET_COOKIE   = "Set-Cookie".freeze
@@ -97,10 +103,10 @@ module ActionDispatch # :nodoc:
 
       def body
         @str_body ||= begin
-                        buf = ''
-                        each { |chunk| buf << chunk }
-                        buf
-                      end
+          buf = ''
+            each { |chunk| buf << chunk }
+          buf
+        end
       end
 
       def write(string)
@@ -112,10 +118,13 @@ module ActionDispatch # :nodoc:
       end
 
       def each(&block)
-        @response.sending!
-        x = @buf.each(&block)
-        @response.sent!
-        x
+        if @str_body
+          return enum_for(:each) unless block_given?
+
+          yield @str_body
+        else
+          each_chunk(&block)
+        end
       end
 
       def abort
@@ -129,6 +138,12 @@ module ActionDispatch # :nodoc:
       def closed?
         @closed
       end
+
+      private
+
+        def each_chunk(&block)
+          @buf.each(&block) # extract into own method
+        end
     end
 
     def self.create(status = 200, header = {}, body = [], default_headers: self.default_headers)