fix: add tests for parallel branches with triggers

Author: cristipufu

tests/runtime/test_resumable.py

@@ -20,8 +20,11 @@
 from uipath_langchain.runtime.storage import SqliteResumableStorage
 
 
-class MockTriggerHandler:
-    """Mock implementation of UiPathResumeTriggerHandler."""
+class SequentialTriggerHandler:
+    """Mock implementation that fires triggers sequentially.
+
+    Resolves triggers one at a time across multiple resume calls.
+    """
 
     def __init__(self):
         self.call_count = 0
@@ -59,9 +62,31 @@ async def read_trigger(self, trigger: UiPathResumeTrigger) -> Any:
         )
 
 
+class ParallelTriggerHandler:
+    """Mock implementation that fires all triggers immediately.
+
+    Resolves all triggers on the first resume call.
+    """
+
+    async def create_trigger(self, suspend_value: Any) -> UiPathResumeTrigger:
+        """Create a trigger from suspend value."""
+        trigger = UiPathResumeTrigger(
+            trigger_type=UiPathResumeTriggerType.API,
+            trigger_name=UiPathResumeTriggerName.API,
+            payload=suspend_value,
+        )
+        return trigger
+
+    async def read_trigger(self, trigger: UiPathResumeTrigger) -> Any:
+        """Read trigger and return immediate response."""
+        assert trigger.payload is not None
+        branch_name = trigger.payload.get("message", "unknown")
+        return f"Response for {branch_name}"
+
+
 @pytest.mark.asyncio
-async def test_parallel_branches_with_multiple_interrupts_execution():
-    """Test graph execution with parallel branches and multiple interrupts."""
+async def test_parallel_branches_with_sequential_trigger_resolution():
+    """Test graph execution with parallel branches where triggers resolve sequentially."""
 
     # Define state
     class State(TypedDict, total=False):
@@ -110,19 +135,19 @@ def branch_c(state: State) -> State:
             # Create base runtime
             base_runtime = UiPathLangGraphRuntime(
                 graph=compiled_graph,
-                runtime_id="parallel-test",
+                runtime_id="parallel-sequential-test",
                 entrypoint="test",
             )
 
             # Create storage and trigger manager
             storage = SqliteResumableStorage(memory)
 
-            # Wrap with UiPathResumableRuntime
+            # Wrap with UiPathResumableRuntime using sequential trigger handler
             runtime = UiPathResumableRuntime(
                 delegate=base_runtime,
                 storage=storage,
-                trigger_manager=MockTriggerHandler(),
-                runtime_id="parallel-test",
+                trigger_manager=SequentialTriggerHandler(),
+                runtime_id="parallel-sequential-test",
             )
 
             # First execution - should hit all 3 interrupts
@@ -141,11 +166,11 @@ def branch_c(state: State) -> State:
             assert len(result.triggers) == 3
 
             # Verify triggers were saved to storage
-            saved_triggers = await storage.get_triggers("parallel-test")
+            saved_triggers = await storage.get_triggers("parallel-sequential-test")
             assert saved_triggers is not None
             assert len(saved_triggers) == 3
 
-            # Resume 1: Resolve only first interrupt (no input, will restore from storage)
+            # Resume 1: Resolve only first interrupt
             result_1 = await runtime.execute(
                 input=None,
                 options=UiPathExecuteOptions(resume=True),
@@ -157,7 +182,7 @@ def branch_c(state: State) -> State:
             assert len(result_1.triggers) == 2
 
             # Verify only 2 triggers remain in storage
-            saved_triggers = await storage.get_triggers("parallel-test")
+            saved_triggers = await storage.get_triggers("parallel-sequential-test")
             assert saved_triggers is not None
             assert len(saved_triggers) == 2
 
@@ -173,7 +198,7 @@ def branch_c(state: State) -> State:
             assert len(result_2.triggers) == 1
 
             # Verify only 1 trigger remains in storage
-            saved_triggers = await storage.get_triggers("parallel-test")
+            saved_triggers = await storage.get_triggers("parallel-sequential-test")
             assert saved_triggers is not None
             assert len(saved_triggers) == 1
 
@@ -188,7 +213,7 @@ def branch_c(state: State) -> State:
             assert result_3.output is not None
 
             # Verify no triggers remain
-            saved_triggers = await storage.get_triggers("parallel-test")
+            saved_triggers = await storage.get_triggers("parallel-sequential-test")
             assert saved_triggers is None or len(saved_triggers) == 0
 
             # Verify all branches completed
@@ -200,3 +225,282 @@ def branch_c(state: State) -> State:
     finally:
         if os.path.exists(temp_db.name):
             os.remove(temp_db.name)
+
+
+@pytest.mark.asyncio
+async def test_parallel_branches_with_parallel_trigger_resolution():
+    """Test graph execution with parallel branches where all triggers fire immediately."""
+
+    # Define state
+    class State(TypedDict, total=False):
+        branch_a_result: str | None
+        branch_b_result: str | None
+        branch_c_result: str | None
+
+    # Define nodes that interrupt
+    def branch_a(state: State) -> State:
+        result = interrupt({"message": "Branch A needs input"})
+        return {"branch_a_result": f"A completed with: {result}"}
+
+    def branch_b(state: State) -> State:
+        result = interrupt({"message": "Branch B needs input"})
+        return {"branch_b_result": f"B completed with: {result}"}
+
+    def branch_c(state: State) -> State:
+        result = interrupt({"message": "Branch C needs input"})
+        return {"branch_c_result": f"C completed with: {result}"}
+
+    # Build graph with parallel branches
+    graph = StateGraph(State)
+    graph.add_node("branch_a", branch_a)
+    graph.add_node("branch_b", branch_b)
+    graph.add_node("branch_c", branch_c)
+
+    # All branches start in parallel
+    graph.add_edge(START, "branch_a")
+    graph.add_edge(START, "branch_b")
+    graph.add_edge(START, "branch_c")
+
+    # All branches go to end
+    graph.add_edge("branch_a", END)
+    graph.add_edge("branch_b", END)
+    graph.add_edge("branch_c", END)
+
+    # Create temporary database
+    temp_db = tempfile.NamedTemporaryFile(delete=False, suffix=".db")
+    temp_db.close()
+
+    try:
+        # Compile graph with checkpointer
+        async with AsyncSqliteSaver.from_conn_string(temp_db.name) as memory:
+            compiled_graph = graph.compile(checkpointer=memory)
+
+            # Create base runtime
+            base_runtime = UiPathLangGraphRuntime(
+                graph=compiled_graph,
+                runtime_id="parallel-parallel-test",
+                entrypoint="test",
+            )
+
+            # Create storage and trigger manager
+            storage = SqliteResumableStorage(memory)
+
+            # Wrap with UiPathResumableRuntime using parallel trigger handler
+            runtime = UiPathResumableRuntime(
+                delegate=base_runtime,
+                storage=storage,
+                trigger_manager=ParallelTriggerHandler(),
+                runtime_id="parallel-parallel-test",
+            )
+
+            # First execution - should hit all 3 interrupts
+            result = await runtime.execute(
+                input={
+                    "branch_a_result": None,
+                    "branch_b_result": None,
+                    "branch_c_result": None,
+                },
+                options=UiPathExecuteOptions(resume=False),
+            )
+
+            # Should be suspended with 3 triggers
+            assert result.status == UiPathRuntimeStatus.SUSPENDED
+            assert result.triggers is not None
+            assert len(result.triggers) == 3
+
+            # Verify triggers were saved to storage
+            saved_triggers = await storage.get_triggers("parallel-parallel-test")
+            assert saved_triggers is not None
+            assert len(saved_triggers) == 3
+
+            # Resume: All triggers should resolve immediately
+            result_resume = await runtime.execute(
+                input=None,
+                options=UiPathExecuteOptions(resume=True),
+            )
+
+            # Should now be successful (all triggers resolved in one go)
+            assert result_resume.status == UiPathRuntimeStatus.SUCCESSFUL
+            assert result_resume.output is not None
+
+            # Verify no triggers remain
+            saved_triggers = await storage.get_triggers("parallel-parallel-test")
+            assert saved_triggers is None or len(saved_triggers) == 0
+
+            # Verify all branches completed
+            output = result_resume.output
+            assert "branch_a_result" in output
+            assert "branch_b_result" in output
+            assert "branch_c_result" in output
+
+            # Verify all branches got their responses
+            assert "Response for Branch A needs input" in output["branch_a_result"]
+            assert "Response for Branch B needs input" in output["branch_b_result"]
+            assert "Response for Branch C needs input" in output["branch_c_result"]
+
+    finally:
+        if os.path.exists(temp_db.name):
+            os.remove(temp_db.name)
+
+
+@pytest.mark.asyncio
+async def test_two_branches_with_two_sequential_interrupts_each():
+    """Test graph execution with 2 parallel branches, each having 2 sequential interrupts."""
+
+    # Define state
+    class State(TypedDict, total=False):
+        branch_a_first_result: str | None
+        branch_a_second_result: str | None
+        branch_b_first_result: str | None
+        branch_b_second_result: str | None
+
+    # Define nodes that interrupt twice sequentially
+    def branch_a_first(state: State) -> State:
+        result = interrupt({"message": "Branch A - First interrupt"})
+        return {"branch_a_first_result": f"A-1 completed with: {result}"}
+
+    def branch_a_second(state: State) -> State:
+        result = interrupt({"message": "Branch A - Second interrupt"})
+        return {"branch_a_second_result": f"A-2 completed with: {result}"}
+
+    def branch_b_first(state: State) -> State:
+        result = interrupt({"message": "Branch B - First interrupt"})
+        return {"branch_b_first_result": f"B-1 completed with: {result}"}
+
+    def branch_b_second(state: State) -> State:
+        result = interrupt({"message": "Branch B - Second interrupt"})
+        return {"branch_b_second_result": f"B-2 completed with: {result}"}
+
+    # Build graph with parallel branches, each with sequential nodes
+    graph = StateGraph(State)
+    graph.add_node("branch_a_first", branch_a_first)
+    graph.add_node("branch_a_second", branch_a_second)
+    graph.add_node("branch_b_first", branch_b_first)
+    graph.add_node("branch_b_second", branch_b_second)
+
+    # Branch A: START -> a_first -> a_second -> END
+    graph.add_edge(START, "branch_a_first")
+    graph.add_edge("branch_a_first", "branch_a_second")
+    graph.add_edge("branch_a_second", END)
+
+    # Branch B: START -> b_first -> b_second -> END
+    graph.add_edge(START, "branch_b_first")
+    graph.add_edge("branch_b_first", "branch_b_second")
+    graph.add_edge("branch_b_second", END)
+
+    # Create temporary database
+    temp_db = tempfile.NamedTemporaryFile(delete=False, suffix=".db")
+    temp_db.close()
+
+    try:
+        # Compile graph with checkpointer
+        async with AsyncSqliteSaver.from_conn_string(temp_db.name) as memory:
+            compiled_graph = graph.compile(checkpointer=memory)
+
+            # Create base runtime
+            base_runtime = UiPathLangGraphRuntime(
+                graph=compiled_graph,
+                runtime_id="two-branches-sequential-test",
+                entrypoint="test",
+            )
+
+            # Create storage and trigger manager
+            storage = SqliteResumableStorage(memory)
+
+            # Wrap with UiPathResumableRuntime using parallel trigger handler
+            runtime = UiPathResumableRuntime(
+                delegate=base_runtime,
+                storage=storage,
+                trigger_manager=ParallelTriggerHandler(),
+                runtime_id="two-branches-sequential-test",
+            )
+
+            # First execution - should hit first interrupt in both branches (2 total)
+            result = await runtime.execute(
+                input={
+                    "branch_a_first_result": None,
+                    "branch_a_second_result": None,
+                    "branch_b_first_result": None,
+                    "branch_b_second_result": None,
+                },
+                options=UiPathExecuteOptions(resume=False),
+            )
+
+            # Should be suspended with 2 triggers (first interrupt from each branch)
+            assert result.status == UiPathRuntimeStatus.SUSPENDED
+            assert result.triggers is not None
+            assert len(result.triggers) == 2
+
+            # Verify triggers were saved to storage
+            saved_triggers = await storage.get_triggers("two-branches-sequential-test")
+            assert saved_triggers is not None
+            assert len(saved_triggers) == 2
+
+            # Verify we got the first interrupts from both branches
+            trigger_messages = [t.payload.get("message") for t in result.triggers]
+            assert "Branch A - First interrupt" in trigger_messages
+            assert "Branch B - First interrupt" in trigger_messages
+
+            # Resume 1: Resolve first interrupts, will hit second interrupts
+            result_1 = await runtime.execute(
+                input=None,
+                options=UiPathExecuteOptions(resume=True),
+            )
+
+            # Should still be suspended with 2 triggers (second interrupt from each branch)
+            assert result_1.status == UiPathRuntimeStatus.SUSPENDED
+            assert result_1.triggers is not None
+            assert len(result_1.triggers) == 2
+
+            # Verify we got the second interrupts from both branches
+            trigger_messages = [t.payload.get("message") for t in result_1.triggers]
+            assert "Branch A - Second interrupt" in trigger_messages
+            assert "Branch B - Second interrupt" in trigger_messages
+
+            # Verify 2 triggers remain in storage
+            saved_triggers = await storage.get_triggers("two-branches-sequential-test")
+            assert saved_triggers is not None
+            assert len(saved_triggers) == 2
+
+            # Resume 2: Resolve second interrupts, should complete
+            result_2 = await runtime.execute(
+                input=None,
+                options=UiPathExecuteOptions(resume=True),
+            )
+
+            # Should now be successful
+            assert result_2.status == UiPathRuntimeStatus.SUCCESSFUL
+            assert result_2.output is not None
+
+            # Verify no triggers remain
+            saved_triggers = await storage.get_triggers("two-branches-sequential-test")
+            assert saved_triggers is None or len(saved_triggers) == 0
+
+            # Verify all branch steps completed
+            output = result_2.output
+            assert "branch_a_first_result" in output
+            assert "branch_a_second_result" in output
+            assert "branch_b_first_result" in output
+            assert "branch_b_second_result" in output
+
+            # Verify all steps got their responses
+            assert (
+                "Response for Branch A - First interrupt"
+                in output["branch_a_first_result"]
+            )
+            assert (
+                "Response for Branch A - Second interrupt"
+                in output["branch_a_second_result"]
+            )
+            assert (
+                "Response for Branch B - First interrupt"
+                in output["branch_b_first_result"]
+            )
+            assert (
+                "Response for Branch B - Second interrupt"
+                in output["branch_b_second_result"]
+            )
+
+    finally:
+        if os.path.exists(temp_db.name):
+            os.remove(temp_db.name)