Add unit tests for _closeables reset in _init for flow reuse

Tests verify that _closeables is properly reset when flows are reused:
- test_flow_reuse_resets_closeables: Base Flow class resets to []
- test_aggregate_by_key_reuse_resets_closeables: AggregateByKey sets [table]
- test_map_with_state_reuse_resets_closeables: MapWithState sets [state] for Table
- test_map_with_state_no_closeables_without_close_method: No closeables for dict state
- test_nosql_target_reuse_resets_closeables: NoSqlTarget sets [table]

Author: Alex Toker

tests/test_flow.py

@@ -5636,3 +5636,174 @@ def test_flow_reuse_with_cycle():
 
         controller.terminate()
         controller.await_termination()
+
+
+def test_flow_reuse_resets_closeables():
+    """Test that _closeables is properly reset when a flow is reused.
+
+    When a flow is run multiple times, _init() should reset _closeables to an empty
+    list to prevent accumulation of closeable resources across runs.
+    """
+    source = SyncEmitSource()
+    map_step = Map(lambda x: x + 1)
+    reduce_step = Reduce(0, lambda acc, x: acc + x)
+
+    source.to(map_step).to(reduce_step)
+
+    # Run the flow multiple times and verify _closeables is reset each time
+    for run_num in range(3):
+        # Before running, _closeables should be empty (from previous _init or initial state)
+        controller = source.run()
+
+        # After run(), _init() has been called and _closeables should be reset
+        # The base Flow class initializes _closeables to [] in _init()
+        assert map_step._closeables == [], f"Run {run_num}: map_step._closeables should be empty"
+        assert reduce_step._closeables == [], f"Run {run_num}: reduce_step._closeables should be empty"
+
+        for i in range(5):
+            controller.emit(i)
+        controller.terminate()
+        result = controller.await_termination()
+        assert result == 15, f"Run {run_num}: Expected 15 but got {result}"
+
+
+def test_aggregate_by_key_reuse_resets_closeables():
+    """Test that AggregateByKey properly resets _closeables on flow reuse.
+
+    AggregateByKey sets _closeables = [self._table] in _init(), which should
+    happen fresh on each run rather than accumulating tables.
+    """
+    from storey import AggregateByKey, FieldAggregator
+    from storey.dtypes import FixedWindows
+
+    table = Table("test", NoopDriver())
+    aggregator = AggregateByKey(
+        [FieldAggregator("sum_col1", "col1", ["sum"], FixedWindows(["1h"]))],
+        table,
+        time_field="time",
+    )
+
+    source = SyncEmitSource()
+    reduce_step = Reduce([], lambda acc, x: acc + [x])
+
+    source.to(aggregator).to(reduce_step)
+
+    base_time = datetime(2020, 7, 21, 12, 0, 0)
+
+    for run_num in range(3):
+        controller = source.run()
+
+        # After _init(), _closeables should contain exactly one table reference
+        assert len(aggregator._closeables) == 1, (
+            f"Run {run_num}: AggregateByKey._closeables should have exactly 1 item, "
+            f"got {len(aggregator._closeables)}"
+        )
+        assert aggregator._closeables[0] is table, f"Run {run_num}: AggregateByKey._closeables[0] should be the table"
+
+        # Emit some data
+        for i in range(3):
+            data = {"col1": i, "time": base_time}
+            controller.emit(data, f"key{i}")
+
+        controller.terminate()
+        controller.await_termination()
+
+
+def test_map_with_state_reuse_resets_closeables():
+    """Test that MapWithState properly resets _closeables on flow reuse.
+
+    MapWithState (via _FunctionWithStateFlow) sets _closeables = [self._state]
+    in _init() when state has a close method (like Table).
+    """
+    table = Table("test_state", NoopDriver())
+
+    def state_fn(event, state):
+        state["count"] = state.get("count", 0) + 1
+        event["count"] = state["count"]
+        return event, state
+
+    source = SyncEmitSource()
+    map_with_state = MapWithState(table, state_fn, group_by_key=True)
+    reduce_step = Reduce([], lambda acc, x: acc + [x])
+
+    source.to(map_with_state).to(reduce_step)
+
+    for run_num in range(3):
+        controller = source.run()
+
+        # Table has a close method, so _closeables should contain exactly the table
+        assert len(map_with_state._closeables) == 1, (
+            f"Run {run_num}: MapWithState._closeables should have exactly 1 item, "
+            f"got {len(map_with_state._closeables)}"
+        )
+        assert map_with_state._closeables[0] is table, f"Run {run_num}: MapWithState._closeables[0] should be the table"
+
+        # Emit some data
+        for i in range(3):
+            controller.emit(Event(body={"value": i}, key=f"key{i}"))
+
+        controller.terminate()
+        controller.await_termination()
+
+
+def test_map_with_state_no_closeables_without_close_method():
+    """Test that MapWithState doesn't add state to _closeables if it has no close method.
+
+    When using a plain dict as state (no close method), _closeables should remain empty.
+    """
+    initial_state = {"count": 0}
+
+    def state_fn(event, state):
+        state["count"] += 1
+        return event["value"] * state["count"], state
+
+    source = SyncEmitSource()
+    map_with_state = MapWithState(initial_state, state_fn, group_by_key=False)
+    reduce_step = Reduce(0, lambda acc, x: acc + x)
+
+    source.to(map_with_state).to(reduce_step)
+
+    for run_num in range(3):
+        controller = source.run()
+
+        # Dict has no close method, so _closeables should be empty
+        assert (
+            map_with_state._closeables == []
+        ), f"Run {run_num}: MapWithState._closeables should be empty when state has no close method"
+
+        for i in range(3):
+            controller.emit({"value": i + 1})
+
+        controller.terminate()
+        controller.await_termination()
+
+
+def test_nosql_target_reuse_resets_closeables():
+    """Test that NoSqlTarget properly resets _closeables on flow reuse.
+
+    NoSqlTarget sets _closeables = [self._table] in _init(), which should
+    happen fresh on each run.
+    """
+    table = Table("test_nosql", NoopDriver())
+
+    source = SyncEmitSource()
+    nosql_target = NoSqlTarget(table)
+
+    source.to(nosql_target)
+
+    for run_num in range(3):
+        controller = source.run()
+
+        # After _init(), _closeables should contain exactly one table reference
+        assert len(nosql_target._closeables) == 1, (
+            f"Run {run_num}: NoSqlTarget._closeables should have exactly 1 item, "
+            f"got {len(nosql_target._closeables)}"
+        )
+        assert nosql_target._closeables[0] is table, f"Run {run_num}: NoSqlTarget._closeables[0] should be the table"
+
+        # Emit some data with keys
+        for i in range(3):
+            controller.emit(Event(body={"col": i}, key=f"key{i}"))
+
+        controller.terminate()
+        controller.await_termination()