From 680fa3c226780ec51c08326a4b64197c7adb724e Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 18 Sep 2024 08:53:47 -0700 Subject: [PATCH] Add one more unit test for exception propagation - we already have plenty of tests for this, but adding one more with a slight variation --- libs/langgraph/tests/test_pregel_async.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 2e031e2f5..5cea88b1d 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -223,6 +223,25 @@ async def iambad(input: Any) -> None: assert inner_task_cancelled +async def test_node_cancellation_on_other_node_exception_two() -> None: + async def awhile(input: Any) -> None: + await asyncio.sleep(1) + + async def iambad(input: Any) -> None: + raise ValueError("I am bad") + + builder = Graph() + builder.add_node("agent", awhile) + builder.add_node("bad", iambad) + builder.set_conditional_entry_point(lambda _: ["agent", "bad"], then=END) + + graph = builder.compile() + + with pytest.raises(ValueError, match="I am bad"): + # This will raise ValueError, not CancelledError + await graph.ainvoke(1) + + @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_dynamic_interrupt(checkpointer_name: str) -> None: class State(TypedDict):