diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index faa359862..319ceed1c 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -82,6 +82,8 @@ from tests.messages import ( _AnyIdToolMessage, ) +pytestmark = pytest.mark.anyio + logger = logging.getLogger(__name__) @@ -5740,34 +5742,6 @@ def test_entrypoint_without_checkpointer() -> None: assert foo.invoke({"a": "1"}, config) == {"current": {"a": "1"}, "previous": None} -async def test_async_entrypoint_without_checkpointer() -> None: - """Test no checkpointer.""" - states = [] - config = {"configurable": {"thread_id": "1"}} - - # Test without previous - @entrypoint() - async def foo(inputs: Any) -> Any: - states.append(inputs) - return inputs - - assert (await foo.ainvoke({"a": "1"}, config)) == {"a": "1"} - - @entrypoint() - async def foo(inputs: Any, *, previous: Any) -> Any: - states.append(previous) - return {"previous": previous, "current": inputs} - - assert (await foo.ainvoke({"a": "1"}, config)) == { - "current": {"a": "1"}, - "previous": None, - } - assert (await foo.ainvoke({"a": "1"}, config)) == { - "current": {"a": "1"}, - "previous": None, - } - - def test_entrypoint_stateful() -> None: """Test stateful entrypoint invoke.""" @@ -5859,26 +5833,6 @@ def test_entrypoint_request_stream_writer() -> None: ] -async def test_entrypoint_from_async_generator() -> None: - """@entrypoint does not support sync generators.""" - # Test invoke - previous_return_values = [] - - # In this version reducers do not work - @entrypoint(checkpointer=MemorySaver()) - async def foo(inputs, previous=None) -> Any: - previous_return_values.append(previous) - yield "a" - yield "b" - - config = {"configurable": {"thread_id": "1"}} - - assert list(await foo.ainvoke({"a": "1"}, config)) == ["a", "b"] - assert previous_return_values == [None] - assert list(foo.invoke({"a": "2"}, config)) == ["a", "b"] - assert previous_return_values == [None, ["a", "b"]] - - @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_multiple_subgraphs( request: pytest.FixtureRequest, checkpointer_name: str diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 7b53870d8..35d6143dd 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -7277,3 +7277,51 @@ async def test_multiple_subgraphs_mixed_checkpointer( ), ((), {"parent_node": {"parent_counter": 7}}), ] + + +@NEEDS_CONTEXTVARS +async def test_async_entrypoint_without_checkpointer() -> None: + """Test no checkpointer.""" + states = [] + config = {"configurable": {"thread_id": "1"}} + + # Test without previous + @entrypoint() + async def foo(inputs: Any) -> Any: + states.append(inputs) + return inputs + + assert (await foo.ainvoke({"a": "1"}, config)) == {"a": "1"} + + @entrypoint() + async def foo(inputs: Any, *, previous: Any) -> Any: + states.append(previous) + return {"previous": previous, "current": inputs} + + assert (await foo.ainvoke({"a": "1"}, config)) == { + "current": {"a": "1"}, + "previous": None, + } + assert (await foo.ainvoke({"a": "1"}, config)) == { + "current": {"a": "1"}, + "previous": None, + } + + +@NEEDS_CONTEXTVARS +async def test_entrypoint_from_async_generator() -> None: + """@entrypoint does not support sync generators.""" + # Test invoke + previous_return_values = [] + + # In this version reducers do not work + @entrypoint(checkpointer=MemorySaver()) + async def foo(inputs, previous=None) -> Any: + previous_return_values.append(previous) + yield "a" + yield "b" + + config = {"configurable": {"thread_id": "1"}} + + assert list(await foo.ainvoke({"a": "1"}, config)) == ["a", "b"] + assert previous_return_values == [None]