Fix style

Author: gtopper

storey/steps/collector.py

@@ -53,7 +53,7 @@ async def _do(self, event):
         # Handle StreamCompletion sentinel
         if isinstance(event, StreamCompletion):
             stream_id = event.original_event.id
-            stream_data = self._collected_streams[stream_id]  # Use [] to trigger defaultdict
+            stream_data = self._collected_streams[stream_id]
 
             stream_data["completions"] += 1
 

tests/test_streaming.py

@@ -1008,7 +1008,8 @@ def stream_chunks(x):
 
         # Collector should receive chunks from both branches and emit collected list
         # Each branch processes each chunk, so we get 4 items total (2 chunks x 2 branches)
-        assert [set(*result)] == [{"a_test_chunk_0", "a_test_chunk_1", "b_test_chunk_0", "b_test_chunk_1"}]
+        assert set(result[0]) == {"a_test_chunk_0", "a_test_chunk_1", "b_test_chunk_0", "b_test_chunk_1"}
+        assert len(result) == 1
 
     def test_async_streaming_graph_split_collector_expected_completions_2(self):
         """Async version: Test streaming through a split with Collector(expected_completions=2)."""
@@ -1036,7 +1037,8 @@ def stream_chunks(x):
             await controller.terminate()
             result = await controller.await_termination()
 
-            assert [set(*result)] == [{"a_test_chunk_0", "a_test_chunk_1", "b_test_chunk_0", "b_test_chunk_1"}]
+            assert set(result[0]) == {"a_test_chunk_0", "a_test_chunk_1", "b_test_chunk_0", "b_test_chunk_1"}
+            assert len(result) == 1
 
         asyncio.run(_test())
 
@@ -1070,7 +1072,7 @@ def stream_chunks(x):
             chunks = list(result)
 
             # Should have 4 chunks total (2 chunks x 2 branches)
-            assert [set(chunks)] == [{"a_test_chunk_0", "a_test_chunk_1", "b_test_chunk_0", "b_test_chunk_1"}]
+            assert set(chunks) == {"a_test_chunk_0", "a_test_chunk_1", "b_test_chunk_0", "b_test_chunk_1"}
         finally:
             controller.terminate()
             controller.await_termination()
@@ -1102,7 +1104,7 @@ def stream_chunks(x):
                 assert inspect.isasyncgen(result)
                 chunks = [chunk async for chunk in result]
 
-                assert [set(chunks)] == [{"a_test_chunk_0", "a_test_chunk_1", "b_test_chunk_0", "b_test_chunk_1"}]
+                assert set(chunks) == {"a_test_chunk_0", "a_test_chunk_1", "b_test_chunk_0", "b_test_chunk_1"}
             finally:
                 await controller.terminate()
                 await controller.await_termination()
@@ -1138,7 +1140,7 @@ def stream_chunks(x):
             chunks = list(result)
 
             # Should have 4 chunks total (2 chunks x 2 branches)
-            assert [set(chunks)] == [{"a_test_chunk_0", "a_test_chunk_1", "b_test_chunk_0", "b_test_chunk_1"}]
+            assert set(chunks) == {"a_test_chunk_0", "a_test_chunk_1", "b_test_chunk_0", "b_test_chunk_1"}
         finally:
             controller.terminate()
             controller.await_termination()
@@ -1169,7 +1171,7 @@ def stream_chunks(x):
                 assert inspect.isasyncgen(result)
                 chunks = [chunk async for chunk in result]
 
-                assert [set(chunks)] == [{"a_test_chunk_0", "a_test_chunk_1", "b_test_chunk_0", "b_test_chunk_1"}]
+                assert set(chunks) == {"a_test_chunk_0", "a_test_chunk_1", "b_test_chunk_0", "b_test_chunk_1"}
             finally:
                 await controller.terminate()
                 await controller.await_termination()