Add async test, add test for validation error

tests/test_pregel.py

@@ -2970,7 +2970,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch(
 def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class(
     checkpoint_at: CheckpointAt,
 ) -> None:
-    from langchain_core.pydantic_v1 import BaseModel
+    from langchain_core.pydantic_v1 import BaseModel, ValidationError
 
     def sorted_add(
         x: list[str], y: Union[list[str], list[tuple[str, str]]]
@@ -3060,6 +3060,9 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class(
                 +---------+                      """
     )
 
+    with pytest.raises(ValidationError):
+        app.invoke({"query": {}})
+
     assert app.invoke({"query": "what is weather in sf"}) == {
         "query": "analyzed: query: what is weather in sf",
         "docs": ["doc1", "doc2", "doc3", "doc4"],

tests/test_pregel_async.py

@@ -2767,6 +2767,146 @@ async def test_in_one_fan_out_state_graph_waiting_edge_via_branch(
     ]
 
 
+@pytest.mark.parametrize(
+    "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
+)
+async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class(
+    checkpoint_at: CheckpointAt,
+) -> None:
+    from langchain_core.pydantic_v1 import BaseModel, ValidationError
+
+    def sorted_add(
+        x: list[str], y: Union[list[str], list[tuple[str, str]]]
+    ) -> list[str]:
+        if isinstance(y[0], tuple):
+            for rem, _ in y:
+                x.remove(rem)
+            y = [t[1] for t in y]
+        return sorted(operator.add(x, y))
+
+    class State(BaseModel):
+        query: str
+        answer: Optional[str] = None
+        docs: Annotated[list[str], sorted_add]
+
+    async def rewrite_query(data: State) -> State:
+        return {"query": f"query: {data.query}"}
+
+    async def analyzer_one(data: State) -> State:
+        return {"query": f"analyzed: {data.query}"}
+
+    async def retriever_one(data: State) -> State:
+        return {"docs": ["doc1", "doc2"]}
+
+    async def retriever_two(data: State) -> State:
+        return {"docs": ["doc3", "doc4"]}
+
+    async def qa(data: State) -> State:
+        return {"answer": ",".join(data.docs)}
+
+    async def decider(data: State) -> str:
+        assert isinstance(data, State)
+        return "retriever_two"
+
+    workflow = StateGraph(State)
+
+    workflow.add_node("rewrite_query", rewrite_query)
+    workflow.add_node("analyzer_one", analyzer_one)
+    workflow.add_node("retriever_one", retriever_one)
+    workflow.add_node("retriever_two", retriever_two)
+    workflow.add_node("qa", qa)
+
+    workflow.set_entry_point("rewrite_query")
+    workflow.add_edge("rewrite_query", "analyzer_one")
+    workflow.add_edge("analyzer_one", "retriever_one")
+    workflow.add_conditional_edges(
+        "rewrite_query", decider, {"retriever_two": "retriever_two"}
+    )
+    workflow.add_edge(["retriever_one", "retriever_two"], "qa")
+    workflow.set_finish_point("qa")
+
+    app = workflow.compile()
+
+    assert app.get_graph().draw_ascii() == (
+        """               +-----------+                     
+               | __start__ |                     
+               +-----------+                     
+                      *                          
+                      *                          
+                      *                          
+             +---------------+                   
+             | rewrite_query |                   
+             +---------------+                   
+               **           **                   
+             **               **                 
+           **                   **               
++--------------+      +-----------------------+  
+| analyzer_one |      | rewrite_query_decider |  
++--------------+      +-----------------------+  
+        *                          *             
+        *                          *             
+        *                          *             
++---------------+         +---------------+      
+| retriever_one |         | retriever_two |      
++---------------+         +---------------+      
+               **           **                   
+                 **       **                     
+                   **   **                       
+                   +----+                        
+                   | qa |                        
+                   +----+                        
+                      *                          
+                      *                          
+                      *                          
+                +---------+                      
+                | __end__ |                      
+                +---------+                      """
+    )
+
+    with pytest.raises(ValidationError):
+        await app.ainvoke({"query": {}})
+
+    assert await app.ainvoke({"query": "what is weather in sf"}) == {
+        "query": "analyzed: query: what is weather in sf",
+        "docs": ["doc1", "doc2", "doc3", "doc4"],
+        "answer": "doc1,doc2,doc3,doc4",
+    }
+
+    assert [c async for c in app.astream({"query": "what is weather in sf"})] == [
+        {"rewrite_query": {"query": "query: what is weather in sf"}},
+        {
+            "analyzer_one": {"query": "analyzed: query: what is weather in sf"},
+            "retriever_two": {"docs": ["doc3", "doc4"]},
+        },
+        {"retriever_one": {"docs": ["doc1", "doc2"]}},
+        {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
+    ]
+
+    app_w_interrupt = workflow.compile(
+        checkpointer=MemorySaverAssertImmutable(at=checkpoint_at),
+        interrupt_after=["retriever_one"],
+    )
+    config = {"configurable": {"thread_id": "1"}}
+
+    assert [
+        c
+        async for c in app_w_interrupt.astream(
+            {"query": "what is weather in sf"}, config
+        )
+    ] == [
+        {"rewrite_query": {"query": "query: what is weather in sf"}},
+        {
+            "analyzer_one": {"query": "analyzed: query: what is weather in sf"},
+            "retriever_two": {"docs": ["doc3", "doc4"]},
+        },
+        {"retriever_one": {"docs": ["doc1", "doc2"]}},
+    ]
+
+    assert [c async for c in app_w_interrupt.astream(None, config)] == [
+        {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
+    ]
+
+
 @pytest.mark.parametrize(
     "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]
 )