Add more tests, ensure we never yield outside the fiber context and that we swap buffers when moving from parent to child.

3 files changed, 18 insertions, 3 deletions

diff --git a/actionpack/lib/action_view/flows.rb b/actionpack/lib/action_view/flows.rb
index b81a34002a..b5e2c5d37d 100644
--- a/actionpack/lib/action_view/flows.rb
+++ b/actionpack/lib/action_view/flows.rb
@@ -22,8 +22,11 @@ module ActionView
   end
 
   class StreamingFlow < OutputFlow
-    def initialize(flow, fiber)
-      @content = flow.content
+    def initialize(view, fiber)
+      @view    = view
+      @parent  = nil
+      @child   = view.output_buffer
+      @content = view._view_flow.content
       @fiber   = fiber
       @root    = Fiber.current.object_id
     end
@@ -36,11 +39,15 @@ module ActionView
       return super if @content.key?(key)
 
       if inside_fiber?
+        view = @view
+
         begin
           @waiting_for = key
+          view.output_buffer, @parent = @child, view.output_buffer
           Fiber.yield
         ensure
           @waiting_for = nil
+          view.output_buffer, @child = @parent, view.output_buffer
         end
       end
 
diff --git a/actionpack/lib/action_view/renderer/streaming_template_renderer.rb b/actionpack/lib/action_view/renderer/streaming_template_renderer.rb
index 1ea00d15ea..5ef6191d79 100644
--- a/actionpack/lib/action_view/renderer/streaming_template_renderer.rb
+++ b/actionpack/lib/action_view/renderer/streaming_template_renderer.rb
@@ -55,7 +55,7 @@ module ActionView
         # Set the view flow to support streaming. It will be aware
         # when to stop rendering the layout because it needs to search
         # something in the template and vice-versa.
-        view._view_flow = StreamingFlow.new(view._view_flow, fiber)
+        view._view_flow = StreamingFlow.new(view, fiber)
 
         # Yo! Start the fiber!
         fiber.resume
diff --git a/actionpack/test/template/streaming_render_test.rb b/actionpack/test/template/streaming_render_test.rb
index 06669bb1ee..754ede9701 100644
--- a/actionpack/test/template/streaming_render_test.rb
+++ b/actionpack/test/template/streaming_render_test.rb
@@ -86,4 +86,12 @@ class FiberedTest < ActiveSupport::TestCase
   def test_render_with_handler_without_streaming_support
     assert_match "<p>This is grand!</p>", buffered_render(:template => "test/hello")
   end
+
+  def test_render_with_streaming_multiple_yields_provide_and_content_for
+    assert_equal "Yes, \nthis works\n like a charm.", buffered_render(:template => "test/streaming", :layout => "layouts/streaming")
+  end
+
+  def test_render_with_streaming_with_fake_yields_and_streaming_buster
+    assert_equal "This won't look\n good.", buffered_render(:template => "test/streaming_buster", :layout => "layouts/streaming")
+  end
 end if defined?(Fiber)
\ No newline at end of file