Minor fixes to the cyclic pr (#598)

Author: davesh0812

storey/flow.py

@@ -44,7 +44,7 @@ def __init__(
         recovery_step=None,
         termination_result_fn=lambda x, y: x if x is not None else y,
         context=None,
-        max_iteration: Optional[int] = None,
+        max_iterations: Optional[int] = None,
         **kwargs,
     ):
         self._outlets = []
@@ -67,7 +67,7 @@ def __init__(
         self._full_event = kwargs.get("full_event")
         self._input_path = kwargs.get("input_path")
         self._result_path = kwargs.get("result_path")
-        self._max_iteration = max_iteration
+        self._max_iterations = max_iterations
         self._runnable = False
         name = kwargs.get("name", None)
         if name:
@@ -284,17 +284,11 @@ def _event_string(event):
     def _should_terminate(self):
         return self._termination_received == len(self._inlets)
 
-    async def _do_downstream(self, event, outlets=None):
-        if not outlets:
-            if event is _termination_obj:
-                outlets = self._outlets
-            else:
-                if self._selected_outlets:
-                    outlet_names = self._selected_outlets
-                    self._selected_outlets = None
-                else:
-                    outlet_names = self.select_outlets(event.body)
-                outlets = self._check_outlets_by_names(outlet_names) if outlet_names else self._outlets
+    async def _do_downstream(self, event, outlets=None, select_outlets: bool = True):
+        if not outlets and event is not _termination_obj and select_outlets:
+            outlet_names = self.select_outlets(event.body)
+            outlets = self._check_outlets_by_names(outlet_names) if outlet_names else None
+        outlets = self._outlets if outlets is None else outlets
 
         if not outlets:
             return
@@ -390,9 +384,9 @@ def _check_step_in_flow(self, type_to_check, visited=None):
         return False
 
     def check_and_update_iteration_number(self, event) -> Optional[Callable]:
-        if hasattr(event, "_cyclic_counter") and self._max_iteration is not None:
+        if hasattr(event, "_cyclic_counter") and self._max_iterations is not None:
             counter = self.get_iteration_counter(event)
-            if counter >= self._max_iteration:
+            if counter >= self._max_iterations:
                 raise RuntimeError(f"Max iterations exceeded in step '{self.name}' for event {event.id}")
             event._cyclic_counter[self.name] = counter + 1
         else:
@@ -457,7 +451,7 @@ def _init(self):
 
     async def _do(self, event):
         if event is _termination_obj:
-            return await self._do_downstream(_termination_obj)
+            return await self._do_downstream(_termination_obj, select_outlets=False)
         else:
             event_body = event if self._full_event else event.body
             outlet_names = self.select_outlets(event_body)
@@ -466,8 +460,8 @@ async def _do(self, event):
                 outlet = self._name_to_outlet["dataframe"]
                 outlets.append(outlet)
             else:
-                outlets = self._check_outlets_by_names(outlet_names) if outlet_names else self._outlets
-            return await self._do_downstream(event, outlets=outlets)
+                outlets = self._check_outlets_by_names(outlet_names)
+            return await self._do_downstream(event, outlets=outlets, select_outlets=False)
 
 
 class Recover(Flow):

tests/test_flow.py

@@ -5538,10 +5538,10 @@ def select(event):
 def test_cyclic_graphs(iterations, with_recovery):
     source = SyncEmitSource()
     my_loop = MyLoop(
-        fn=lambda x: x, iterations=iterations, name="my_loop", end="end", counter="counter", max_iteration=5
+        fn=lambda x: x, iterations=iterations, name="my_loop", end="end", counter="counter", max_iterations=5
     )
     start = Map(lambda x: x, name="start")
-    counter = Map(lambda x: x + 1, name="counter", max_iteration=5)
+    counter = Map(lambda x: x + 1, name="counter", max_iterations=5)
     end = Map(lambda x: x, name="end")
 
     source.to(start)