ML-1142: Fix completion on error and reject emit after termination. (#…

…296)

storey/flow.py

@@ -145,6 +145,10 @@ async def _do_and_recover(self, event):
             if recovery_step is None:
                 if self.context and hasattr(self.context, 'push_error'):
                     message = traceback.format_exc()
+                    if event._awaitable_result:
+                        none_or_coroutine = event._awaitable_result._set_error(ex)
+                        if none_or_coroutine:
+                            await none_or_coroutine
                     if self.logger:
                         self.logger.error(f'Pushing error to error stream: {ex}\n{message}')
                     self.context.push_error(event, f"{ex}\n{message}", source=self.name)

storey/sources.py

@@ -197,6 +197,9 @@ def __init__(self, buffer_size: Optional[int] = None, key_field: Union[list, str
         self._ex = None
         self._closeables = []
 
+    def _init(self):
+        self._is_terminated = False
+
     async def _run_loop(self):
         loop = asyncio.get_running_loop()
         self._termination_future = asyncio.get_running_loop().create_future()
@@ -239,6 +242,10 @@ def _raise_on_error(self, ex):
     def _emit(self, event):
         if event is not _termination_obj:
             self._raise_on_error(self._ex)
+            if self._is_terminated:
+                raise ValueError('Cannot emit to a terminated flow')
+        else:
+            self._is_terminated = True
         self._q.put(event)
         if event is not _termination_obj:
             self._raise_on_error(self._ex)
@@ -368,6 +375,9 @@ def __init__(self, buffer_size: int = 1024, key_field: Union[list, str, None] =
         self._ex = None
         self._closeables = []
 
+    def _init(self):
+        self._is_terminated = False
+
     async def _run_loop(self):
         while True:
             event = await self._q.get()
@@ -401,6 +411,10 @@ def _raise_on_error(self):
     async def _emit(self, event):
         if event is not _termination_obj:
             self._raise_on_error()
+            if self._is_terminated:
+                raise ValueError('Cannot emit to a terminated flow')
+        else:
+            self._is_terminated = True
         await self._q.put(event)
         if event is not _termination_obj:
             self._raise_on_error()

tests/test_flow.py

@@ -944,7 +944,7 @@ def boom(_):
     for i in range(2):
         try:
             awaitable_results.append(controller.emit(0))
-        except ValueError as ex:
+        except ValueError:
             pass
 
     last_trace_size = None