From b99a101b2c0c807ac2ed13bb0e426c7f64e46e2f Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Tue, 3 Sep 2024 15:34:11 -0700
Subject: [PATCH] Add large parallel graph test

---
 libs/langgraph/tests/test_pregel_async.py | 68 ++++++++++++++++++++++-
 1 file changed, 67 insertions(+), 1 deletion(-)

diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py
index dd89901b8..7707e2405 100644
--- a/libs/langgraph/tests/test_pregel_async.py
+++ b/libs/langgraph/tests/test_pregel_async.py
@@ -1,6 +1,7 @@
 import asyncio
 import json
 import operator
+import random
 import re
 import sys
 from collections import Counter
@@ -514,7 +515,6 @@ async def test_cancel_graph_astream_events_v2(checkpointer_name: Optional[str])
                 if chunk["event"] == "on_chain_stream" and not chunk["parent_ids"]:
                     got_event = True
                     assert chunk["data"]["chunk"] == {"alittlewhile": {"value": 2}}
-                    await asyncio.sleep(0.1)
                     break
 
         # did break
@@ -8431,6 +8431,72 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
         ]
 
 
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
+async def test_large_graph(checkpointer_name: str) -> None:
+    class OverallState(TypedDict):
+        subjects: list[str]
+        jokes: Annotated[list[str], operator.add]
+
+    async def continue_to_jokes(state: OverallState):
+        return [Send("generate_joke", {"subject": s}) for s in state["subjects"]]
+
+    class JokeInput(TypedDict):
+        subject: str
+
+    class JokeOutput(TypedDict):
+        jokes: list[str]
+
+    async def edit(state: JokeInput):
+        subject = state["subject"]
+        return {"subject": f"{subject} - hohoho"}
+
+    # subgraph
+    subgraph = StateGraph(input=JokeInput, output=JokeOutput)
+    subgraph.add_node("edit", edit)
+    subgraph.add_node(
+        "generate", lambda state: {"jokes": [f"Joke about {state['subject']}"]}
+    )
+    subgraph.add_node(
+        "bump", lambda state: {"jokes": [state["jokes"][0] + " a"]}, input=JokeOutput
+    )
+    subgraph.set_entry_point("edit")
+    subgraph.add_edge("edit", "generate")
+    subgraph.add_edge("generate", "bump")
+    subgraph.add_conditional_edges(
+        "bump", lambda state: END if state["jokes"][0].endswith(" a" * 10) else "bump"
+    )
+    subgraph.set_finish_point("generate")
+
+    # parent graph
+    builder = StateGraph(OverallState)
+    builder.add_node("generate_joke", subgraph.compile())
+    builder.add_conditional_edges(START, continue_to_jokes)
+    builder.add_edge("generate_joke", END)
+
+    async with awith_checkpointer(checkpointer_name) as checkpointer:
+        graph = builder.compile(checkpointer=checkpointer)
+        config = {"configurable": {"thread_id": "1"}}
+
+        # invoke and pause at nested interrupt
+        assert (
+            len(
+                [
+                    c
+                    async for c in graph.astream(
+                        {
+                            "subjects": [
+                                random.choice("abcdefghijklmnopqrstuvwxyz")
+                                for _ in range(1000)
+                            ]
+                        },
+                        config=config,
+                    )
+                ]
+            )
+            == 1000
+        )
+
+
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
 async def test_send_to_nested_graphs(checkpointer_name: str) -> None:
     class OverallState(TypedDict):