test: Ensure that async checkpointers in tests are setup in same event loop as they are used in

- This more closely remembers the environment they're used in, so it's what we should be testing
- Remove unnecessary pytest-asyncio dependency, use anyio pytest plugin instead
- Convert remaining async tests using only memory checkpointer to use all existing ones
This commit is contained in:
Nuno Campos
2024-09-02 17:24:54 -07:00
parent a937754132
commit d65d0ec665
14 changed files with 5975 additions and 4414 deletions
+1 -1
View File
@@ -15,7 +15,7 @@ coverage:
--cov-report term-missing:skip-covered
start-postgres:
docker compose -f tests/compose-postgres.yml up -V --force-recreate --wait
docker compose -f tests/compose-postgres.yml up -V --force-recreate --wait --remove-orphans
stop-postgres:
docker compose -f tests/compose-postgres.yml down -v
+8 -27
View File
@@ -2829,13 +2829,13 @@ diagrams = ["jinja2", "railroad-diagrams"]
[[package]]
name = "pytest"
version = "7.4.4"
version = "8.3.2"
description = "pytest: simple powerful testing with Python"
optional = false
python-versions = ">=3.7"
python-versions = ">=3.8"
files = [
{file = "pytest-7.4.4-py3-none-any.whl", hash = "sha256:b090cdf5ed60bf4c45261be03239c2c1c22df034fbffe691abe93cd80cea01d8"},
{file = "pytest-7.4.4.tar.gz", hash = "sha256:2cf0005922c6ace4a3e2ec8b4080eb0d9753fdc93107415332f50ce9e7994280"},
{file = "pytest-8.3.2-py3-none-any.whl", hash = "sha256:4ba08f9ae7dcf84ded419494d229b48d0903ea6407b030eaec46df5e6a73bba5"},
{file = "pytest-8.3.2.tar.gz", hash = "sha256:c132345d12ce551242c87269de812483f5bcc87cdbb4722e48487ba194f9fdce"},
]
[package.dependencies]
@@ -2843,29 +2843,11 @@ colorama = {version = "*", markers = "sys_platform == \"win32\""}
exceptiongroup = {version = ">=1.0.0rc8", markers = "python_version < \"3.11\""}
iniconfig = "*"
packaging = "*"
pluggy = ">=0.12,<2.0"
tomli = {version = ">=1.0.0", markers = "python_version < \"3.11\""}
pluggy = ">=1.5,<2"
tomli = {version = ">=1", markers = "python_version < \"3.11\""}
[package.extras]
testing = ["argcomplete", "attrs (>=19.2.0)", "hypothesis (>=3.56)", "mock", "nose", "pygments (>=2.7.2)", "requests", "setuptools", "xmlschema"]
[[package]]
name = "pytest-asyncio"
version = "0.20.3"
description = "Pytest support for asyncio"
optional = false
python-versions = ">=3.7"
files = [
{file = "pytest-asyncio-0.20.3.tar.gz", hash = "sha256:83cbf01169ce3e8eb71c6c278ccb0574d1a7a3bb8eaaf5e50e0ad342afb33b36"},
{file = "pytest_asyncio-0.20.3-py3-none-any.whl", hash = "sha256:f129998b209d04fcc65c96fc85c11e5316738358909a8399e93be553d7656442"},
]
[package.dependencies]
pytest = ">=6.1.0"
[package.extras]
docs = ["sphinx (>=5.3)", "sphinx-rtd-theme (>=1.0)"]
testing = ["coverage (>=6.2)", "flaky (>=3.5.0)", "hypothesis (>=5.7.1)", "mypy (>=0.931)", "pytest-trio (>=0.7.0)"]
dev = ["argcomplete", "attrs (>=19.2)", "hypothesis (>=3.56)", "mock", "pygments (>=2.7.2)", "requests", "setuptools", "xmlschema"]
[[package]]
name = "pytest-cov"
@@ -3069,7 +3051,6 @@ files = [
{file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"},
{file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"},
{file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"},
{file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"},
{file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"},
{file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"},
{file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"},
@@ -4310,4 +4291,4 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools",
[metadata]
lock-version = "2.0"
python-versions = ">=3.9.0,<4.0"
content-hash = "7e0d6fd967987fcf46b038c47b5ca81d86cb8fff89704eac2074f2ce55336a2b"
content-hash = "e0787721e6cb80996a39284c74f1ddee4789c6060d6c2272cb1957f4686bd478"
+1 -3
View File
@@ -14,10 +14,9 @@ langgraph-checkpoint = "^1.0.2"
[tool.poetry.group.dev.dependencies]
pytest = "^7.3.0"
pytest = "^8.3.2"
pytest-cov = "^4.0.0"
pytest-dotenv = "^0.5.2"
pytest-asyncio = "^0.20.3"
pytest-mock = "^3.10.0"
syrupy = "^4.0.2"
httpx = "^0.26.0"
@@ -73,7 +72,6 @@ requires = ["poetry-core>=1.0.0"]
build-backend = "poetry.core.masonry.api"
[tool.pytest.ini_options]
asyncio_mode = "auto"
# --strict-markers will raise errors on unknown marks.
# https://docs.pytest.org/en/7.1.x/how-to/mark.html#raising-errors-on-unknown-marks
#
File diff suppressed because it is too large Load Diff
+42 -63
View File
@@ -1,8 +1,6 @@
import asyncio
import sys
from concurrent.futures import ThreadPoolExecutor
from contextlib import asynccontextmanager, contextmanager
from typing import AsyncIterator, Iterator, TypeVar
from contextlib import asynccontextmanager
from typing import AsyncIterator, Optional
from uuid import UUID, uuid4
import pytest
@@ -10,6 +8,7 @@ from psycopg import AsyncConnection, Connection
from psycopg_pool import AsyncConnectionPool, ConnectionPool
from pytest_mock import MockerFixture
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.checkpoint.postgres import PostgresSaver
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from langgraph.checkpoint.sqlite import SqliteSaver
@@ -19,6 +18,11 @@ from tests.memory_assert import MemorySaverAssertImmutable
DEFAULT_POSTGRES_URI = "postgres://postgres:postgres@localhost:5442/"
@pytest.fixture
def anyio_backend():
return "asyncio"
@pytest.fixture()
def deterministic_uuids(mocker: MockerFixture) -> MockerFixture:
side_effect = (
@@ -27,35 +31,6 @@ def deterministic_uuids(mocker: MockerFixture) -> MockerFixture:
return mocker.patch("uuid.uuid4", side_effect=side_effect)
"""
pytest-asyncio doesn't support calling async fixtures with getfixturevalue
so we need to use ThreadPoolExecutor to run the async fixture in a thread
https://github.com/pytest-dev/pytest-asyncio/issues/112#issuecomment-462062890
"""
T = TypeVar("T")
def close_loop(loop: asyncio.AbstractEventLoop) -> None:
loop.run_until_complete(loop.shutdown_asyncgens())
loop.run_until_complete(loop.shutdown_default_executor())
asyncio.set_event_loop(None)
loop.close()
@contextmanager
def agen_to_gen(agen: AsyncIterator[T]) -> Iterator[T]:
with ThreadPoolExecutor(1) as bg:
loop = asyncio.new_event_loop()
bg.submit(asyncio.set_event_loop, loop).result()
try:
yield bg.submit(loop.run_until_complete, agen.__aenter__()).result()
finally:
bg.submit(
loop.run_until_complete, agen.__aexit__(None, None, None)
).result()
bg.submit(close_loop, loop).result()
# checkpointer fixtures
@@ -70,12 +45,6 @@ def checkpointer_sqlite():
yield checkpointer
@pytest.fixture(scope="function")
def checkpointer_sqlite_aio():
with agen_to_gen(_checkpointer_sqlite_aio()) as checkpointer:
yield checkpointer
@asynccontextmanager
async def _checkpointer_sqlite_aio():
async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer:
@@ -143,16 +112,10 @@ def checkpointer_postgres_pool():
conn.execute(f"DROP DATABASE {database}")
@pytest.fixture(scope="function")
def checkpointer_postgres_aio():
if sys.version_info < (3, 10):
pytest.skip("Async Postgres tests require Python 3.10+")
with agen_to_gen(_checkpointer_postgres_aio()) as checkpointer:
yield checkpointer
@asynccontextmanager
async def _checkpointer_postgres_aio():
if sys.version_info < (3, 10):
pytest.skip("Async Postgres tests require Python 3.10+")
database = f"test_{uuid4().hex[:16]}"
# create unique db
async with await AsyncConnection.connect(
@@ -174,16 +137,10 @@ async def _checkpointer_postgres_aio():
await conn.execute(f"DROP DATABASE {database}")
@pytest.fixture(scope="function")
def checkpointer_postgres_aio_pipe():
if sys.version_info < (3, 10):
pytest.skip("Async Postgres tests require Python 3.10+")
with agen_to_gen(_checkpointer_postgres_aio_pipe()) as checkpointer:
yield checkpointer
@asynccontextmanager
async def _checkpointer_postgres_aio_pipe():
if sys.version_info < (3, 10):
pytest.skip("Async Postgres tests require Python 3.10+")
database = f"test_{uuid4().hex[:16]}"
# create unique db
async with await AsyncConnection.connect(
@@ -208,16 +165,10 @@ async def _checkpointer_postgres_aio_pipe():
await conn.execute(f"DROP DATABASE {database}")
@pytest.fixture(scope="function")
def checkpointer_postgres_aio_pool():
if sys.version_info < (3, 10):
pytest.skip("Async Postgres tests require Python 3.10+")
with agen_to_gen(_checkpointer_postgres_aio_pool()) as checkpointer:
yield checkpointer
@asynccontextmanager
async def _checkpointer_postgres_aio_pool():
if sys.version_info < (3, 10):
pytest.skip("Async Postgres tests require Python 3.10+")
database = f"test_{uuid4().hex[:16]}"
# create unique db
async with await AsyncConnection.connect(
@@ -240,6 +191,30 @@ async def _checkpointer_postgres_aio_pool():
await conn.execute(f"DROP DATABASE {database}")
@asynccontextmanager
async def awith_checkpointer(
checkpointer_name: Optional[str],
) -> AsyncIterator[BaseCheckpointSaver]:
if checkpointer_name is None:
yield None
elif checkpointer_name == "memory":
yield MemorySaverAssertImmutable()
elif checkpointer_name == "sqlite_aio":
async with _checkpointer_sqlite_aio() as checkpointer:
yield checkpointer
elif checkpointer_name == "postgres_aio":
async with _checkpointer_postgres_aio() as checkpointer:
yield checkpointer
elif checkpointer_name == "postgres_aio_pipe":
async with _checkpointer_postgres_aio_pipe() as checkpointer:
yield checkpointer
elif checkpointer_name == "postgres_aio_pool":
async with _checkpointer_postgres_aio_pool() as checkpointer:
yield checkpointer
else:
raise NotImplementedError(f"Unknown checkpointer: {checkpointer_name}")
ALL_CHECKPOINTERS_SYNC = [
"memory",
"sqlite",
@@ -254,3 +229,7 @@ ALL_CHECKPOINTERS_ASYNC = [
"postgres_aio_pipe",
"postgres_aio_pool",
]
ALL_CHECKPOINTERS_ASYNC_PLUS_NONE = [
*ALL_CHECKPOINTERS_ASYNC,
None,
]
+2
View File
@@ -8,6 +8,8 @@ from langgraph.channels.last_value import LastValue
from langgraph.channels.topic import Topic
from langgraph.errors import EmptyChannelError, InvalidUpdateError
pytestmark = pytest.mark.anyio
def test_last_value() -> None:
with LastValue(int).from_checkpoint(None, {}) as channel:
+21 -20
View File
@@ -4,12 +4,16 @@ import pytest
from pytest_mock import MockerFixture
from langgraph.graph import END, START, StateGraph
@pytest.mark.parametrize(
"checkpointer_name",
["memory", "sqlite", "postgres", "postgres_pipe"],
from tests.conftest import (
ALL_CHECKPOINTERS_ASYNC,
ALL_CHECKPOINTERS_SYNC,
awith_checkpointer,
)
pytestmark = pytest.mark.anyio
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
def test_interruption_without_state_updates(
request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture
) -> None:
@@ -47,12 +51,9 @@ def test_interruption_without_state_updates(
assert graph.get_state(thread).next == ()
@pytest.mark.parametrize(
"checkpointer_name",
["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"],
)
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_interruption_without_state_updates_async(
request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture
checkpointer_name: str, mocker: MockerFixture
):
"""Test interruption without state updates. This test confirms that
interrupting doesn't require a state key having been updated in the prev step"""
@@ -72,17 +73,17 @@ async def test_interruption_without_state_updates_async(
builder.add_edge("step_2", "step_3")
builder.add_edge("step_3", END)
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
graph = builder.compile(checkpointer=checkpointer, interrupt_after="*")
async with awith_checkpointer(checkpointer_name) as checkpointer:
graph = builder.compile(checkpointer=checkpointer, interrupt_after="*")
initial_input = {"input": "hello world"}
thread = {"configurable": {"thread_id": "1"}}
initial_input = {"input": "hello world"}
thread = {"configurable": {"thread_id": "1"}}
await graph.ainvoke(initial_input, thread, debug=True)
assert (await graph.aget_state(thread)).next == ("step_2",)
await graph.ainvoke(initial_input, thread, debug=True)
assert (await graph.aget_state(thread)).next == ("step_2",)
await graph.ainvoke(None, thread, debug=True)
assert (await graph.aget_state(thread)).next == ("step_3",)
await graph.ainvoke(None, thread, debug=True)
assert (await graph.aget_state(thread)).next == ("step_3",)
await graph.ainvoke(None, thread, debug=True)
assert (await graph.aget_state(thread)).next == ()
await graph.ainvoke(None, thread, debug=True)
assert (await graph.aget_state(thread)).next == ()
+35 -39
View File
@@ -23,8 +23,15 @@ from pydantic import BaseModel as BaseModelV2
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.prebuilt import ToolNode, ValidationNode, create_react_agent
from langgraph.prebuilt.tool_node import InjectedState
from tests.conftest import (
ALL_CHECKPOINTERS_ASYNC,
ALL_CHECKPOINTERS_SYNC,
awith_checkpointer,
)
from tests.messages import _AnyIdHumanMessage
pytestmark = pytest.mark.anyio
class FakeToolCallingModel(BaseChatModel):
def _generate(
@@ -53,10 +60,7 @@ class FakeToolCallingModel(BaseChatModel):
return self
@pytest.mark.parametrize(
"checkpointer_name",
["memory", "sqlite", "postgres", "postgres_pipe"],
)
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
def test_no_modifier(request: pytest.FixtureRequest, checkpointer_name: str) -> None:
checkpointer: BaseCheckpointSaver = request.getfixturevalue(
"checkpointer_" + checkpointer_name
@@ -89,43 +93,35 @@ def test_no_modifier(request: pytest.FixtureRequest, checkpointer_name: str) ->
assert saved.pending_writes == []
@pytest.mark.parametrize(
"checkpointer_name",
["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"],
)
async def test_no_modifier_async(
request: pytest.FixtureRequest, checkpointer_name: str
) -> None:
checkpointer: BaseCheckpointSaver = request.getfixturevalue(
f"checkpointer_{checkpointer_name}"
)
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_no_modifier_async(checkpointer_name: str) -> None:
async with awith_checkpointer(checkpointer_name) as checkpointer:
model = FakeToolCallingModel()
model = FakeToolCallingModel()
agent = create_react_agent(model, [], checkpointer=checkpointer)
inputs = [HumanMessage("hi?")]
thread = {"configurable": {"thread_id": "123"}}
response = await agent.ainvoke({"messages": inputs}, thread, debug=True)
expected_response = {"messages": inputs + [AIMessage(content="hi?", id="0")]}
assert response == expected_response
agent = create_react_agent(model, [], checkpointer=checkpointer)
inputs = [HumanMessage("hi?")]
thread = {"configurable": {"thread_id": "123"}}
response = await agent.ainvoke({"messages": inputs}, thread, debug=True)
expected_response = {"messages": inputs + [AIMessage(content="hi?", id="0")]}
assert response == expected_response
if checkpointer:
saved = await checkpointer.aget_tuple(thread)
assert saved is not None
assert saved.checkpoint["channel_values"] == {
"messages": [
_AnyIdHumanMessage(content="hi?"),
AIMessage(content="hi?", id="0"),
],
"agent": "agent",
}
assert saved.metadata == {
"parents": {},
"source": "loop",
"writes": {"agent": {"messages": [AIMessage(content="hi?", id="0")]}},
"step": 1,
}
assert saved.pending_writes == []
if checkpointer:
saved = await checkpointer.aget_tuple(thread)
assert saved is not None
assert saved.checkpoint["channel_values"] == {
"messages": [
_AnyIdHumanMessage(content="hi?"),
AIMessage(content="hi?", id="0"),
],
"agent": "agent",
}
assert saved.metadata == {
"parents": {},
"source": "loop",
"writes": {"agent": {"messages": [AIMessage(content="hi?", id="0")]}},
"step": 1,
}
assert saved.pending_writes == []
def test_passing_two_modifiers():
+1 -27
View File
@@ -74,10 +74,7 @@ from langgraph.store.memory import MemoryStore
from tests.any_str import AnyDict, AnyStr, AnyVersion, UnsortedSequence
from tests.conftest import ALL_CHECKPOINTERS_SYNC
from tests.fake_tracer import FakeTracer
from tests.memory_assert import (
MemorySaverAssertCheckpointMetadata,
MemorySaverNoPending,
)
from tests.memory_assert import MemorySaverAssertCheckpointMetadata
from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage
@@ -1626,29 +1623,6 @@ def test_cond_edge_after_send() -> None:
assert graph.invoke(["0"]) == ["0", "1", "2", "2", "3"]
async def test_checkpointer_null_pending_writes() -> None:
class Node:
def __init__(self, name: str):
self.name = name
setattr(self, "__name__", name)
def __call__(self, state):
return [self.name]
builder = StateGraph(Annotated[list, operator.add])
builder.add_node(Node("1"))
builder.add_edge(START, "1")
graph = builder.compile(checkpointer=MemorySaverNoPending())
assert graph.invoke([], {"configurable": {"thread_id": "foo"}}) == ["1"]
assert graph.invoke([], {"configurable": {"thread_id": "foo"}}) == ["1"] * 2
assert (await graph.ainvoke([], {"configurable": {"thread_id": "foo"}})) == [
"1"
] * 3
assert (await graph.ainvoke([], {"configurable": {"thread_id": "foo"}})) == [
"1"
] * 4
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
def test_invoke_checkpoint_three(
mocker: MockerFixture, request: pytest.FixtureRequest, checkpointer_name: str
File diff suppressed because it is too large Load Diff
+3 -1
View File
@@ -1,3 +1,4 @@
import warnings
from typing import Annotated as Annotated2
from typing import Any
@@ -45,7 +46,8 @@ def test_warns_invalid_schema(schema: Any):
)
def test_doesnt_warn_valid_schema(schema: Any):
# Assert the function does not raise a warning
with pytest.warns(None):
with warnings.catch_warnings():
warnings.simplefilter("error")
_warn_invalid_state_schema(schema)
+3
View File
@@ -1,11 +1,14 @@
import asyncio
from typing import Any, Optional
import pytest
from pytest_mock import MockerFixture
from langgraph.store.base import BaseStore
from langgraph.store.batch import AsyncBatchedStore
pytestmark = pytest.mark.anyio
async def test_async_batch_store(mocker: MockerFixture) -> None:
aget = mocker.stub()
+2
View File
@@ -11,6 +11,8 @@ from langgraph.graph import END, StateGraph
from langgraph.graph.graph import CompiledGraph
from langgraph.utils import is_async_callable, is_async_generator
pytestmark = pytest.mark.anyio
def test_is_async() -> None:
async def func() -> None: