From 64cfbb0d021958b1cd17d0eb24a82186b161c0d9 Mon Sep 17 00:00:00 2001 From: Sydney Runkle Date: Tue, 29 Apr 2025 13:51:45 -0400 Subject: [PATCH] remove v1 test --- libs/langgraph/tests/test_pregel.py | 172 ---------------------------- 1 file changed, 172 deletions(-) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index a365d02f0..97e6786a9 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -2596,178 +2596,6 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch( ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1( - snapshot: SnapshotAssertion, - mocker: MockerFixture, - request: pytest.FixtureRequest, - checkpointer_name: str, -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - setup = mocker.Mock() - teardown = mocker.Mock() - - @contextmanager - def assert_ctx_once() -> Iterator[None]: - assert setup.call_count == 0 - assert teardown.call_count == 0 - try: - yield - finally: - assert setup.call_count == 1 - assert teardown.call_count == 1 - setup.reset_mock() - teardown.reset_mock() - - @contextmanager - def make_httpx_client() -> Iterator[httpx.Client]: - setup() - with httpx.Client() as client: - try: - yield client - finally: - teardown() - - 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 InnerObject(BaseModel): - yo: int - - class State(BaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True) - - query: str - inner: Annotated[InnerObject, lambda x, y: y] - answer: Optional[str] = None - docs: Annotated[list[str], sorted_add] - client: Annotated[httpx.Client, Context(make_httpx_client)] - - class Input(BaseModel): - query: str - inner: InnerObject - - class Output(BaseModel): - answer: str - docs: list[str] - - class StateUpdate(BaseModel): - query: Optional[str] = None - answer: Optional[str] = None - docs: Optional[list[str]] = None - - class UpdateDocs34(BaseModel): - docs: list[str] = ["doc3", "doc4"] - - def rewrite_query(data: State) -> State: - assert isinstance(data.inner, InnerObject) - return {"query": f"query: {data.query}"} - - def analyzer_one(data: State) -> State: - assert isinstance(data.inner, InnerObject) - return StateUpdate(query=f"analyzed: {data.query}") - - def retriever_one(data: State) -> State: - return {"docs": ["doc1", "doc2"]} - - def retriever_two(data: State) -> State: - time.sleep(0.1) - return UpdateDocs34() - - def qa(data: State) -> State: - return {"answer": ",".join(data.docs)} - - def decider(data: State) -> str: - assert isinstance(data, State) - return "retriever_two" - - workflow = StateGraph(State, input=Input, output=Output) - - 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() - - if checkpointer_name == "memory": - assert app.get_graph().draw_mermaid(with_styles=False) == snapshot - assert app.get_input_jsonschema() == snapshot - assert app.get_output_jsonschema() == snapshot - - with pytest.raises(ValidationError), assert_ctx_once(): - app.invoke({"query": {}}) - - with assert_ctx_once(): - assert app.invoke({"query": "what is weather in sf", "inner": {"yo": 1}}) == { - "docs": ["doc1", "doc2", "doc3", "doc4"], - "answer": "doc1,doc2,doc3,doc4", - } - - with assert_ctx_once(): - assert [ - *app.stream({"query": "what is weather in sf", "inner": {"yo": 1}}) - ] == [ - {"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=checkpointer, - interrupt_after=["retriever_one"], - ) - config = {"configurable": {"thread_id": "1"}} - - with assert_ctx_once(): - assert [ - c - for c in app_w_interrupt.stream( - {"query": "what is weather in sf", "inner": {"yo": 1}}, 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"]}}, - {"__interrupt__": ()}, - ] - - with assert_ctx_once(): - assert [c for c in app_w_interrupt.stream(None, config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] - - with assert_ctx_once(): - assert app_w_interrupt.update_state( - config, {"docs": ["doc5"]}, as_node="rewrite_query" - ) == { - "configurable": { - "thread_id": "1", - "checkpoint_id": AnyStr(), - "checkpoint_ns": "", - } - } - - @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( snapshot: SnapshotAssertion,