New tests and related fixes

Author: gtopper

storey/flow.py

@@ -485,7 +485,8 @@ def _init(self):
     async def _do(self, event):
         if event is _termination_obj:
             return await self._do_downstream(_termination_obj, select_outlets=False)
-        # StreamCompletion objects should be forwarded to all outlets without routing
+        # StreamCompletion should propagate to all outlets (like _termination_obj)
+        # to avoid hangs in cyclic graphs and ensure all Collectors receive completions
         if isinstance(event, StreamCompletion):
             return await self._do_downstream(event, select_outlets=False)
         else:

storey/steps/collector.py

@@ -12,9 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+import copy
 from collections import defaultdict
 
-from ..dtypes import Event, StreamCompletion, _termination_obj
+from ..dtypes import StreamCompletion, _termination_obj
 from ..flow import Flow
 
 
@@ -41,7 +42,9 @@ def __init__(self, expected_completions: int = 1, **kwargs):
             raise ValueError("expected_completions must be at least 1")
         self._expected_completions = expected_completions
         # Map from event id -> {"chunks": [], "completions": 0, "first_event": Event}
-        self._collected_streams = defaultdict(lambda: {"chunks": [], "completions": 0, "first_event": None})
+        self._collected_streams: dict[str, dict] = defaultdict(
+            lambda: {"chunks": [], "completions": 0, "first_event": None}
+        )
 
     async def _do(self, event):
         if event is _termination_obj:
@@ -50,51 +53,43 @@ async def _do(self, event):
         # Handle StreamCompletion sentinel
         if isinstance(event, StreamCompletion):
             stream_id = event.original_event.id
-            stream_data = self._collected_streams.get(stream_id)
-
-            if stream_data is None:
-                # No chunks received for this stream - unusual but possible
-                return
+            stream_data = self._collected_streams[stream_id]  # Use [] to trigger defaultdict
 
             stream_data["completions"] += 1
 
             if stream_data["completions"] >= self._expected_completions:
                 # Stream is complete - emit collected result
-                first_event = stream_data["first_event"]
-                if first_event:
-                    collected_body = [chunk.body for chunk in stream_data["chunks"]]
-                    if len(collected_body) == 1:
-                        collected_body = collected_body[0]
+                # Use first_event if we have chunks, otherwise use original_event from completion (empty stream)
+                base_event = stream_data["first_event"] or event.original_event
+                collected_body = [chunk.body for chunk in stream_data["chunks"]]
+                if len(collected_body) == 1:
+                    collected_body = collected_body[0]
 
-                    # Create a new event with the collected body
-                    collected_event = Event(
-                        body=collected_body,
-                        key=first_event.key,
-                        processing_time=first_event.processing_time,
-                        id=first_event.id,
-                        headers=first_event.headers,
-                        method=first_event.method,
-                        path=first_event.path,
-                        content_type=first_event.content_type,
-                        awaitable_result=first_event._awaitable_result,
-                    )
-                    await self._do_downstream(collected_event)
+                # Copy the original event to preserve all attributes (important for offset management)
+                collected_event = copy.copy(base_event)
+                collected_event.body = collected_body
+                # Clear streaming attributes
+                if hasattr(collected_event, "streaming_step"):
+                    del collected_event.streaming_step
+                if hasattr(collected_event, "chunk_id"):
+                    del collected_event.chunk_id
+                await self._do_downstream(collected_event)
 
                 # Clean up
                 del self._collected_streams[stream_id]
-            return
+            return None
 
         # Check if this is a streaming chunk (has streaming_step attribute)
-        streaming_step = getattr(event, "streaming_step", None)
-        if streaming_step:
+        if hasattr(event, "streaming_step"):
             stream_id = event.id
             stream_data = self._collected_streams[stream_id]
             if stream_data["first_event"] is None:
                 stream_data["first_event"] = event
             stream_data["chunks"].append(event)
+            return None
         else:
             # Non-streaming event - pass through directly
-            await self._do_downstream(event)
+            return await self._do_downstream(event)
 
     async def _cleanup(self):
         # Warn about incomplete streams on cleanup

tests/test_flow.py

@@ -5592,7 +5592,7 @@ def test_cyclic_graphs(iterations, with_recovery):
     counter.to(my_loop)
     my_loop.to(end)
     end.to(Complete())
-    my_loop._outlets.append(counter)
+    my_loop.to(counter)
     if with_recovery:
         recovery_step = Map(lambda x: -1, name="end-2")
         counter.set_recovery_step(recovery_step)
@@ -5636,10 +5636,10 @@ def test_two_cyclic_graphs():
     start.to(counter)
     counter.to(my_loop)
     my_loop.to(counter_2)
-    my_loop._outlets.append(counter)
+    my_loop.to(counter)
     counter_2.to(my_loop_2)
     my_loop_2.to(end)
-    my_loop_2._outlets.append(counter_2)
+    my_loop_2.to(counter_2)
     end.to(Complete())
     controller = source.run()
 
@@ -5666,7 +5666,7 @@ def test_flow_reuse_with_cycle():
     my_loop.to(end)
     end.to(Complete())
     # Create the cycle by appending counter as an outlet of my_loop
-    my_loop._outlets.append(counter)
+    my_loop.to(counter)
 
     # Run the SAME flow 3 times to test reusability with cyclic structure
     for run_num in range(3):

tests/test_streaming.py

@@ -458,6 +458,56 @@ def stream_chunks(x):
 
         asyncio.run(_test())
 
+    def test_collector_empty_stream(self):
+        """Test that Collector emits an empty list for a stream with zero chunks."""
+
+        def empty_stream(x):
+            return
+            yield  # Makes it a generator
+
+        controller = build_flow(
+            [
+                SyncEmitSource(),
+                Map(empty_stream),
+                Collector(),
+                Reduce([], lambda acc, x: acc + [x]),
+            ]
+        ).run()
+
+        controller.emit("test")
+        controller.terminate()
+        result = controller.await_termination()
+
+        # Empty stream should emit an empty list
+        assert len(result) == 1
+        assert result[0] == []
+
+    def test_async_collector_empty_stream(self):
+        """Async version: Test that Collector emits an empty list for a stream with zero chunks."""
+
+        async def _test():
+            def empty_stream(x):
+                return
+                yield  # Makes it a generator
+
+            controller = build_flow(
+                [
+                    AsyncEmitSource(),
+                    Map(empty_stream),
+                    Collector(),
+                    Reduce([], lambda acc, x: acc + [x]),
+                ]
+            ).run()
+
+            await controller.emit("test")
+            await controller.terminate()
+            result = await controller.await_termination()
+
+            assert len(result) == 1
+            assert result[0] == []
+
+        asyncio.run(_test())
+
 
 class TestCompleteStreaming:
     """Tests for Complete step streaming support."""
@@ -1190,13 +1240,16 @@ def select_outlets(self, event):
         controller.terminate()
         result = controller.await_termination()
 
-        # Should have 2 collected results (one per emit)
-        assert len(result) == 2
-        # low_value chunks go to branch_low
-        low_result = [r for r in result if any("LOW_" in str(item) for item in r)]
-        high_result = [r for r in result if any("HIGH_" in str(item) for item in r)]
+        # 4 results: chunks go to one branch, but StreamCompletion goes to all branches
+        # (like _termination_obj) to avoid hangs in cyclic graphs.
+        # Branches that don't receive chunks emit empty lists.
+        assert len(result) == 4
+        low_result = [r for r in result if r and any("LOW_" in str(item) for item in r)]
+        high_result = [r for r in result if r and any("HIGH_" in str(item) for item in r)]
+        empty_results = [r for r in result if r == []]
         assert len(low_result) == 1
         assert len(high_result) == 1
+        assert len(empty_results) == 2
         assert "LOW_low_value_chunk_0" in low_result[0]
         assert "HIGH_high_value_chunk_0" in high_result[0]
 
@@ -1235,11 +1288,14 @@ def select_outlets(self, event):
             await controller.terminate()
             result = await controller.await_termination()
 
-            assert len(result) == 2
-            low_result = [r for r in result if any("LOW_" in str(item) for item in r)]
-            high_result = [r for r in result if any("HIGH_" in str(item) for item in r)]
+            # 4 results: chunks go to one branch, but StreamCompletion goes to all branches
+            assert len(result) == 4
+            low_result = [r for r in result if r and any("LOW_" in str(item) for item in r)]
+            high_result = [r for r in result if r and any("HIGH_" in str(item) for item in r)]
+            empty_results = [r for r in result if r == []]
             assert len(low_result) == 1
             assert len(high_result) == 1
+            assert len(empty_results) == 2
             assert "LOW_low_value_chunk_0" in low_result[0]
             assert "HIGH_high_value_chunk_0" in high_result[0]
 
@@ -1403,3 +1459,86 @@ def failing_transform(x):
                     await controller.await_termination()
 
         asyncio.run(_test())
+
+    def test_streaming_in_cycle_fails(self):
+        """Test that streaming step inside a cycle fails on second iteration.
+
+        When a streaming step is inside a cycle, the first iteration streams chunks.
+        When those chunks loop back to the streaming step, they already have
+        streaming_step set, so the step should fail with StreamingError.
+        """
+
+        class AlwaysLoop(Map):
+            """A Map step that always routes back to the loop target."""
+
+            def __init__(self, loop_target, **kwargs):
+                super().__init__(**kwargs)
+                self._loop_target = loop_target
+
+            def select_outlets(self, event_body):
+                # Always loop back - the streaming error should stop us
+                return [self._loop_target]
+
+        def stream_chunks(x):
+            yield f"{x}_chunk_0"
+            yield f"{x}_chunk_1"
+
+        source = SyncEmitSource()
+        # The streaming map is the entry point of the loop
+        streaming_map = Map(stream_chunks, name="streamer", max_iterations=5)
+        loop_controller = AlwaysLoop(
+            fn=lambda x: x, name="loop_ctrl", loop_target="streamer", max_iterations=5
+        )
+        end = Reduce([], lambda acc, x: acc + [x], name="end")
+
+        source.to(streaming_map)
+        streaming_map.to(loop_controller)
+        loop_controller.to(end)
+        loop_controller.to(streaming_map)  # Create cycle
+
+        controller = source.run()
+
+        controller.emit("test")
+        controller.terminate()
+
+        # Should fail because chunks looping back already have streaming_step set
+        with pytest.raises(StreamingError, match="Streaming on top of streaming is not allowed"):
+            controller.await_termination()
+
+    def test_async_streaming_in_cycle_fails(self):
+        """Async version: Test that streaming step inside a cycle fails on second iteration."""
+
+        async def _test():
+            class AlwaysLoop(Map):
+                def __init__(self, loop_target, **kwargs):
+                    super().__init__(**kwargs)
+                    self._loop_target = loop_target
+
+                def select_outlets(self, event_body):
+                    return [self._loop_target]
+
+            def stream_chunks(x):
+                yield f"{x}_chunk_0"
+                yield f"{x}_chunk_1"
+
+            source = AsyncEmitSource()
+            streaming_map = Map(stream_chunks, name="streamer", max_iterations=5)
+            loop_controller = AlwaysLoop(
+                fn=lambda x: x, name="loop_ctrl", loop_target="streamer", max_iterations=5
+            )
+            end = Reduce([], lambda acc, x: acc + [x], name="end")
+
+            source.to(streaming_map)
+            streaming_map.to(loop_controller)
+            loop_controller.to(end)
+            loop_controller.to(streaming_map)  # Create cycle
+
+            controller = source.run()
+
+            await controller.emit("test")
+            await controller.terminate()
+
+            with pytest.raises(StreamingError, match="Streaming on top of streaming is not allowed"):
+                await controller.await_termination()
+
+        asyncio.run(_test())