nicen it up

Author: adrianlyjak

src/workflows/handler.py

@@ -130,4 +130,5 @@ async def cancel_run(self) -> None:
         """
         if self.ctx:
             self.ctx._internal_cancel_run()
-            await asyncio.sleep(0)
+            if self._run_task is not None:
+                await self._run_task

src/workflows/runtime/broker.py

@@ -260,6 +260,7 @@ async def _run_workflow() -> None:
                 self._state.is_running = False
 
                 if exception_raised:
+                    print("exception_raised", exception_raised)
                     # cancel the stream
                     self.write_event_to_stream(StopEvent())
 

src/workflows/server/representation_utils.py

@@ -98,7 +98,7 @@ def _extract_workflow_structure(
     # First pass: Add all nodes
     for step_name, step_func in steps.items():
         step_config = step_func._step_config
-        
+
         # Add step node
         step_label = (
             _truncate_label(step_name, max_label_length)
@@ -195,7 +195,7 @@ def _extract_workflow_structure(
     # Second pass: Add edges
     for step_name, step_func in steps.items():
         step_config = step_func._step_config
-        
+
         # Edges from steps to return types
         for return_type in step_config.return_types:
             if return_type is not type(None):

src/workflows/workflow.py

@@ -352,7 +352,7 @@ def _validate(self) -> bool:
         steps_accepting_stop_event: list[str] = []
 
         for name, step_func in self._get_steps().items():
-            step_config: StepConfig  = step_func._step_config
+            step_config: StepConfig = step_func._step_config
 
             # Check that no user-defined step accepts StopEvent (only _done step should)
             if name != "_done":

tests/test_handler.py

@@ -10,7 +10,8 @@
 from workflows.handler import WorkflowHandler
 
 
-def test_str() -> None:
+@pytest.mark.asyncio
+async def test_str() -> None:
     h = WorkflowHandler()
     h.set_result([])
     assert str(h) == "[]"

tests/test_workflow.py

@@ -75,7 +75,7 @@ async def slow_step(self, ev: StartEvent) -> StopEvent:
             return StopEvent(result="Done")
 
     with pytest.raises(WorkflowTimeoutError):
-        await WorkflowTestRunner(SlowWorkflow(timeout=1)).run()
+        await WorkflowTestRunner(SlowWorkflow(timeout=0.1)).run()
 
 
 @pytest.mark.asyncio
@@ -442,21 +442,31 @@ async def step(self, ctx: Context, ev: StartEvent) -> StopEvent:
 
 @pytest.mark.asyncio
 async def test_workflow_context_to_dict(workflow: Workflow) -> None:
-    handler = workflow.run()
-    ctx = handler.ctx
+    ctx: Union[Context, None] = None
+    new_ctx: Union[Context, None] = None
+    try:
+        handler = workflow.run()
+        ctx = handler.ctx
 
-    ctx.send_event(EventWithName(name="test"))  # type:ignore
+        ctx.send_event(EventWithName(name="test"))  # type:ignore
 
-    # get the context dict
-    data = ctx.to_dict()  # type:ignore
+        # get the context dict
+        data = ctx.to_dict()  # type:ignore
 
-    # finish workflow
-    await handler
+        # finish workflow
+        await handler
 
-    new_ctx = Context.from_dict(workflow, data)
-    new_ctx._init_broker(workflow)
-    assert new_ctx._broker_run is not None
-    assert new_ctx._broker_run._state.queues["start_step"].get_nowait().name == "test"
+        new_ctx = Context.from_dict(workflow, data)
+        new_ctx._init_broker(workflow)
+        assert new_ctx._broker_run is not None
+        assert (
+            new_ctx._broker_run._state.queues["start_step"].get_nowait().name == "test"
+        )
+    finally:
+        if ctx is not None and ctx._broker_run is not None:
+            await ctx._broker_run.shutdown()
+        if new_ctx is not None and new_ctx._broker_run is not None:
+            await new_ctx._broker_run.shutdown()
 
 
 class HumanInTheLoopWorkflow(Workflow):