Add test to confirm only one root run is produced

langchain-ai · Sep 4, 2024 · dc10b3a · dc10b3a
1 parent af780f9
commit dc10b3a
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py
@@ -9823,12 +9823,18 @@ def edit(state: JokeState):
 
     graph = builder.compile(checkpointer=checkpointer)
     config = {"configurable": {"thread_id": "1"}}
+    tracer = FakeTracer()
 
     # invoke and pause at nested interrupt
-    assert graph.invoke({"subjects": ["cats", "dogs"]}, config=config) == {
+    assert graph.invoke(
+        {"subjects": ["cats", "dogs"]}, config={**config, "callbacks": [tracer]}
+    ) == {
         "subjects": ["cats", "dogs"],
         "jokes": [],
     }
+    assert len(tracer.runs) == 1, "Should produce exactly 1 root run"
+
+    # check state
     outer_state = graph.get_state(config)
     assert outer_state == StateSnapshot(
         values={"subjects": ["cats", "dogs"], "jokes": []},

diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py
@@ -8467,15 +8467,19 @@ async def edit(state: JokeState):
     async with awith_checkpointer(checkpointer_name) as checkpointer:
         graph = builder.compile(checkpointer=checkpointer)
         config = {"configurable": {"thread_id": "1"}}
+        tracer = FakeTracer()
 
         # invoke and pause at nested interrupt
-        assert await graph.ainvoke({"subjects": ["cats", "dogs"]}, config=config) == {
+        assert await graph.ainvoke(
+            {"subjects": ["cats", "dogs"]}, config={**config, "callbacks": [tracer]}
+        ) == {
             "subjects": ["cats", "dogs"],
             "jokes": [],
         }
+        assert len(tracer.runs) == 1, "Should produce exactly 1 root run"
+
         # check state
         outer_state = await graph.aget_state(config)
-
         assert outer_state == StateSnapshot(
             values={"subjects": ["cats", "dogs"], "jokes": []},
             tasks=(