Enable async tests that were being skipped

- async tests are placed in test_pregel_async, not in test_pregel
- to avoid tests placed in wrong file being accidentally skipped i've added the auto-async mark to sync test file
This commit is contained in:
Nuno Campos
2025-01-20 11:57:19 -08:00
parent 12be3fac33
commit d48dec5452
2 changed files with 50 additions and 48 deletions
+2 -48
View File
@@ -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
+48
View File
@@ -7272,3 +7272,51 @@ async def test_multiple_subgraphs_mixed_checkpointer(
),
((), {"parent_node": {"parent_counter": 7}}),
]
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,
}
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"]]