From 76e207b643451394074cef3d489aa1fb98c2cdee Mon Sep 17 00:00:00 2001
From: vbarda <vadym@langchain.dev>
Date: Fri, 7 Jun 2024 16:36:48 -0400
Subject: [PATCH] add unittest

---
 tests/__snapshots__/test_pregel.ambr | 61 ++++++++++++++++++++++++++++
 tests/test_pregel.py                 | 42 +++++++++++++++++++
 2 files changed, 103 insertions(+)

diff --git a/tests/__snapshots__/test_pregel.ambr b/tests/__snapshots__/test_pregel.ambr
index 68f8a32b9..7277e8cf4 100644
--- a/tests/__snapshots__/test_pregel.ambr
+++ b/tests/__snapshots__/test_pregel.ambr
@@ -114,6 +114,67 @@
   
   '''
 # ---
+# name: test_conditional_entrypoint_to_multiple_state_graph
+  '{"title": "LangGraphInput", "$ref": "#/definitions/OverallState", "definitions": {"OverallState": {"title": "OverallState", "type": "object", "properties": {"locations": {"title": "Locations", "type": "array", "items": {"type": "string"}}, "results": {"title": "Results", "type": "array", "items": {"type": "string"}}}, "required": ["locations", "results"]}}}'
+# ---
+# name: test_conditional_entrypoint_to_multiple_state_graph.1
+  '{"title": "LangGraphOutput", "type": "object", "properties": {"locations": {"title": "Locations", "type": "array", "items": {"type": "string"}}, "results": {"title": "Results", "type": "array", "items": {"type": "string"}}}}'
+# ---
+# name: test_conditional_entrypoint_to_multiple_state_graph.2
+  '''
+  {
+    "nodes": [
+      {
+        "id": "__start__",
+        "type": "schema",
+        "data": "__start__"
+      },
+      {
+        "id": "__end__",
+        "type": "schema",
+        "data": "__end__"
+      },
+      {
+        "id": "get_weather",
+        "type": "runnable",
+        "data": {
+          "id": [
+            "langgraph",
+            "utils",
+            "RunnableCallable"
+          ],
+          "name": "get_weather"
+        }
+      }
+    ],
+    "edges": [
+      {
+        "source": "get_weather",
+        "target": "__end__"
+      },
+      {
+        "source": "__start__",
+        "target": "get_weather",
+        "conditional": true
+      },
+      {
+        "source": "__start__",
+        "target": "__end__",
+        "conditional": true
+      }
+    ]
+  }
+  '''
+# ---
+# name: test_conditional_entrypoint_to_multiple_state_graph.3
+  '''
+  graph TD;
+  	get_weather --> __end__;
+  	__start__ -.-> get_weather;
+  	__start__ -.-> __end__;
+  
+  '''
+# ---
 # name: test_conditional_entrypoint_graph_state
   '{"title": "LangGraphInput", "$ref": "#/definitions/AgentState", "definitions": {"AgentState": {"title": "AgentState", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "output": {"title": "Output", "type": "string"}, "steps": {"title": "Steps", "type": "array", "items": {"type": "string"}}}}}}'
 # ---
diff --git a/tests/test_pregel.py b/tests/test_pregel.py
index 798d11241..96603dc0c 100644
--- a/tests/test_pregel.py
+++ b/tests/test_pregel.py
@@ -1934,6 +1934,48 @@ def test_conditional_entrypoint_graph(snapshot: SnapshotAssertion) -> None:
     ]
 
 
+def test_conditional_entrypoint_to_multiple_state_graph(
+    snapshot: SnapshotAssertion,
+) -> None:
+    class OverallState(TypedDict):
+        locations: list[str]
+        results: Annotated[list[str], operator.add]
+
+    def get_weather(state: str) -> str:
+        location = state["location"]
+        weather = "sunny" if len(location) > 2 else "cloudy"
+        return {"results": [f"It's {weather} in {location}"]}
+
+    def continue_to_weather(state: OverallState) -> str | list[str]:
+        return [
+            Send("get_weather", {"location": location})
+            for location in state["locations"]
+        ]
+
+    workflow = StateGraph(OverallState)
+
+    workflow.add_node("get_weather", get_weather)
+    workflow.add_edge("get_weather", END)
+    workflow.set_conditional_entry_point(continue_to_weather)
+
+    app = workflow.compile()
+
+    assert app.get_input_schema().schema_json() == snapshot
+    assert app.get_output_schema().schema_json() == snapshot
+    assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
+    assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
+
+    assert app.invoke({"locations": ["sf", "nyc"]}, debug=True) == {
+        "locations": ["sf", "nyc"],
+        "results": ["It's cloudy in sf", "It's sunny in nyc"],
+    }
+
+    assert [*app.stream({"locations": ["sf", "nyc"]}, stream_mode="values")][-1] == {
+        "locations": ["sf", "nyc"],
+        "results": ["It's cloudy in sf", "It's sunny in nyc"],
+    }
+
+
 def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None:
     from langchain_core.agents import AgentAction, AgentFinish
     from langchain_core.language_models.fake import FakeStreamingListLLM