From 92cc3f0e0e0a7a6d3d223d61d5ce3d3aec64f00d Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 9 May 2025 08:38:31 -0700 Subject: [PATCH 1/2] More idiomatic sync/async store fixtures in pytest - Use a single fixture that returns the store, instead of list of names --- libs/langgraph/tests/conftest.py | 58 +++++++++++++---------- libs/langgraph/tests/test_pregel.py | 24 ++++------ libs/langgraph/tests/test_pregel_async.py | 29 ++++++------ 3 files changed, 58 insertions(+), 53 deletions(-) diff --git a/libs/langgraph/tests/conftest.py b/libs/langgraph/tests/conftest.py index 7a4829e03..a6aa7586e 100644 --- a/libs/langgraph/tests/conftest.py +++ b/libs/langgraph/tests/conftest.py @@ -1,6 +1,6 @@ import sys from collections.abc import AsyncIterator, Iterator -from contextlib import asynccontextmanager +from contextlib import asynccontextmanager, contextmanager from typing import Optional from uuid import UUID, uuid4 @@ -374,8 +374,8 @@ def cache(request: pytest.FixtureRequest) -> Iterator[BaseCache]: raise ValueError(f"Unknown cache type: {request.param}") -@pytest.fixture(scope="function") -def store_postgres(): +@contextmanager +def _store_postgres(): database = f"test_{uuid4().hex[:16]}" # create unique db with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: @@ -391,8 +391,8 @@ def store_postgres(): conn.execute(f"DROP DATABASE {database}") -@pytest.fixture(scope="function") -def store_postgres_pipe(): +@contextmanager +def _store_postgres_pipe(): database = f"test_{uuid4().hex[:16]}" # create unique db with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: @@ -411,8 +411,8 @@ def store_postgres_pipe(): conn.execute(f"DROP DATABASE {database}") -@pytest.fixture(scope="function") -def store_postgres_pool(): +@contextmanager +def _store_postgres_pool(): database = f"test_{uuid4().hex[:16]}" # create unique db with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: @@ -430,13 +430,35 @@ def store_postgres_pool(): conn.execute(f"DROP DATABASE {database}") -@pytest.fixture(scope="function") -def store_in_memory(): - yield InMemoryStore() +@pytest.fixture( + scope="function", + params=["in_memory", "postgres", "postgres_pipe", "postgres_pool"], +) +def sync_store(request: pytest.FixtureRequest) -> Iterator[BaseStore]: + store_name = request.param + if store_name is None: + yield None + elif store_name == "in_memory": + yield InMemoryStore() + elif store_name == "postgres": + with _store_postgres() as store: + yield store + elif store_name == "postgres_pipe": + with _store_postgres_pipe() as store: + yield store + elif store_name == "postgres_pool": + with _store_postgres_pool() as store: + yield store + else: + raise NotImplementedError(f"Unknown store {store_name}") -@asynccontextmanager -async def awith_store(store_name: Optional[str]) -> AsyncIterator[BaseStore]: +@pytest.fixture( + scope="function", + params=["in_memory", "postgres_aio", "postgres_aio_pipe", "postgres_aio_pool"], +) +async def async_store(request: pytest.FixtureRequest) -> AsyncIterator[BaseStore]: + store_name = request.param if store_name is None: yield None elif store_name == "in_memory": @@ -483,15 +505,3 @@ ALL_CHECKPOINTERS_ASYNC_PLUS_NONE = [ *ALL_CHECKPOINTERS_ASYNC, None, ] -ALL_STORES_SYNC = [ - "in_memory", - "postgres", - "postgres_pipe", - "postgres_pool", -] -ALL_STORES_ASYNC = [ - "in_memory", - "postgres_aio", - "postgres_aio_pipe", - "postgres_aio_pool", -] diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index d26a6e408..389f11427 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -71,7 +71,6 @@ from tests.agents import AgentAction, AgentFinish from tests.any_str import AnyStr, AnyVersion, FloatBetween, UnsortedSequence from tests.conftest import ( ALL_CHECKPOINTERS_SYNC, - ALL_STORES_SYNC, REGULAR_CHECKPOINTERS_SYNC, SHOULD_CHECK_SNAPSHOTS, ) @@ -5063,12 +5062,10 @@ def test_multiple_sinks_subgraphs(snapshot: SnapshotAssertion) -> None: @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -@pytest.mark.parametrize("store_name", ALL_STORES_SYNC) def test_store_injected( - request: pytest.FixtureRequest, checkpointer_name: str, store_name: str + request: pytest.FixtureRequest, checkpointer_name: str, sync_store: BaseStore ) -> None: checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - the_store = request.getfixturevalue(f"store_{store_name}") class State(TypedDict): count: Annotated[int, operator.add] @@ -5112,7 +5109,7 @@ def test_store_injected( builder.add_node(f"node_{i}", Node(i)) builder.add_edge("__start__", f"node_{i}") - graph = builder.compile(store=the_store, checkpointer=checkpointer) + graph = builder.compile(store=sync_store, checkpointer=checkpointer) results = graph.batch( [{"count": 0}] * M, @@ -5121,25 +5118,25 @@ def test_store_injected( ) result = results[-1] assert result == {"count": N + 1} - returned_doc = the_store.get(namespace, doc_id).value + returned_doc = sync_store.get(namespace, doc_id).value assert returned_doc == {**doc, "from_thread": thread_1, "some_val": 0} - assert len(the_store.search(namespace)) == 1 + assert len(sync_store.search(namespace)) == 1 # Check results after another turn of the same thread result = graph.invoke({"count": 0}, {"configurable": {"thread_id": thread_1}}) assert result == {"count": (N + 1) * 2} - returned_doc = the_store.get(namespace, doc_id).value + returned_doc = sync_store.get(namespace, doc_id).value assert returned_doc == {**doc, "from_thread": thread_1, "some_val": N + 1} - assert len(the_store.search(namespace)) == 1 + assert len(sync_store.search(namespace)) == 1 result = graph.invoke({"count": 0}, {"configurable": {"thread_id": thread_2}}) assert result == {"count": N + 1} - returned_doc = the_store.get(namespace, doc_id).value + returned_doc = sync_store.get(namespace, doc_id).value assert returned_doc == { **doc, "from_thread": thread_2, "some_val": 0, } # Overwrites the whole doc - assert len(the_store.search(namespace)) == 1 # still overwriting the same one + assert len(sync_store.search(namespace)) == 1 # still overwriting the same one def test_enum_node_names(): @@ -7392,16 +7389,15 @@ def test_entrypoint_with_return_and_save() -> None: assert previous_ == ["hello", "goodbye"] -def test_overriding_injectable_args_with_tasks() -> None: +def test_overriding_injectable_args_with_tasks(sync_store: BaseStore) -> None: """Test overriding injectable args in tasks.""" - from langgraph.store.memory import InMemoryStore @task def foo(store: BaseStore, writer: StreamWriter, value: Any) -> None: assert store is value assert writer is value - @entrypoint(store=InMemoryStore()) + @entrypoint(store=sync_store) def main(inputs, store: BaseStore) -> str: assert store is not None foo(store=None, writer=None, value=None).result() diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 0c815f49e..e01ee06fe 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -69,11 +69,9 @@ from tests.any_str import AnyStr, AnyVersion, FloatBetween, UnsortedSequence from tests.conftest import ( ALL_CHECKPOINTERS_ASYNC, ALL_CHECKPOINTERS_ASYNC_PLUS_NONE, - ALL_STORES_ASYNC, REGULAR_CHECKPOINTERS_ASYNC, SHOULD_CHECK_SNAPSHOTS, awith_checkpointer, - awith_store, ) from tests.fake_tracer import FakeTracer from tests.memory_assert import MemorySaverNoPending @@ -6468,8 +6466,9 @@ async def test_checkpointer_null_pending_writes() -> None: @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -@pytest.mark.parametrize("store_name", ALL_STORES_ASYNC) -async def test_store_injected_async(checkpointer_name: str, store_name: str) -> None: +async def test_store_injected_async( + checkpointer_name: str, async_store: BaseStore +) -> None: class State(TypedDict): count: Annotated[int, operator.add] @@ -6527,9 +6526,8 @@ async def test_store_injected_async(checkpointer_name: str, store_name: str) -> async with ( awith_checkpointer(checkpointer_name) as checkpointer, - awith_store(store_name) as the_store, ): - graph = builder.compile(store=the_store, checkpointer=checkpointer) + graph = builder.compile(store=async_store, checkpointer=checkpointer) # Test batch operations with multiple threads results = await graph.abatch( @@ -6539,32 +6537,32 @@ async def test_store_injected_async(checkpointer_name: str, store_name: str) -> ) result = results[-1] assert result == {"count": N + 1} - returned_doc = (await the_store.aget(namespace, doc_id)).value + returned_doc = (await async_store.aget(namespace, doc_id)).value assert returned_doc == {**doc, "from_thread": thread_1, "some_val": 0} - assert len(await the_store.asearch(namespace)) == 1 + assert len(await async_store.asearch(namespace)) == 1 # Check results after another turn of the same thread result = await graph.ainvoke( {"count": 0}, {"configurable": {"thread_id": thread_1}} ) assert result == {"count": (N + 1) * 2} - returned_doc = (await the_store.aget(namespace, doc_id)).value + returned_doc = (await async_store.aget(namespace, doc_id)).value assert returned_doc == {**doc, "from_thread": thread_1, "some_val": N + 1} - assert len(await the_store.asearch(namespace)) == 1 + assert len(await async_store.asearch(namespace)) == 1 # Test with a different thread result = await graph.ainvoke( {"count": 0}, {"configurable": {"thread_id": thread_2}} ) assert result == {"count": N + 1} - returned_doc = (await the_store.aget(namespace, doc_id)).value + returned_doc = (await async_store.aget(namespace, doc_id)).value assert returned_doc == { **doc, "from_thread": thread_2, "some_val": 0, } # Overwrites the whole doc assert ( - len(await the_store.asearch(namespace)) == 1 + len(await async_store.asearch(namespace)) == 1 ) # still overwriting the same one @@ -8220,16 +8218,17 @@ async def test_named_tasks_functional() -> None: @NEEDS_CONTEXTVARS -async def test_overriding_injectable_args_with_async_task() -> None: +async def test_overriding_injectable_args_with_async_task( + async_store: BaseStore, +) -> None: """Test overriding injectable args in tasks.""" - from langgraph.store.memory import InMemoryStore @task async def foo(store: BaseStore, writer: StreamWriter, value: Any) -> None: assert store is value assert writer is value - @entrypoint(store=InMemoryStore()) + @entrypoint(store=async_store) async def main(inputs, store: BaseStore) -> str: assert store is not None await foo(store=None, writer=None, value=None) From dd55e5097b06c6c7a729b8a7510a6a5809b146d8 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 9 May 2025 09:01:25 -0700 Subject: [PATCH 2/2] Move to sep file --- libs/langgraph/tests/conftest.py | 145 +++-------------------- libs/langgraph/tests/conftest_store.py | 154 +++++++++++++++++++++++++ 2 files changed, 168 insertions(+), 131 deletions(-) create mode 100644 libs/langgraph/tests/conftest_store.py diff --git a/libs/langgraph/tests/conftest.py b/libs/langgraph/tests/conftest.py index a6aa7586e..0796112e3 100644 --- a/libs/langgraph/tests/conftest.py +++ b/libs/langgraph/tests/conftest.py @@ -1,6 +1,6 @@ import sys from collections.abc import AsyncIterator, Iterator -from contextlib import asynccontextmanager, contextmanager +from contextlib import asynccontextmanager from typing import Optional from uuid import UUID, uuid4 @@ -24,8 +24,15 @@ from langgraph.checkpoint.serde.encrypted import EncryptedSerializer from langgraph.checkpoint.sqlite import SqliteSaver from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver from langgraph.store.base import BaseStore -from langgraph.store.memory import InMemoryStore -from langgraph.store.postgres import AsyncPostgresStore, PostgresStore +from tests.conftest_store import ( + _store_memory, + _store_postgres, + _store_postgres_aio, + _store_postgres_aio_pipe, + _store_postgres_aio_pool, + _store_postgres_pipe, + _store_postgres_pool, +) pytest.register_assert_rewrite("tests.memory_assert") @@ -294,76 +301,6 @@ async def awith_checkpointer( raise NotImplementedError(f"Unknown checkpointer: {checkpointer_name}") -@asynccontextmanager -async def _store_postgres_aio(): - if sys.version_info < (3, 10): - pytest.skip("Async Postgres tests require Python 3.10+") - database = f"test_{uuid4().hex[:16]}" - async with await AsyncConnection.connect( - DEFAULT_POSTGRES_URI, autocommit=True - ) as conn: - await conn.execute(f"CREATE DATABASE {database}") - try: - async with AsyncPostgresStore.from_conn_string( - DEFAULT_POSTGRES_URI + database - ) as store: - await store.setup() - yield store - finally: - async with await AsyncConnection.connect( - DEFAULT_POSTGRES_URI, autocommit=True - ) as conn: - await conn.execute(f"DROP DATABASE {database}") - - -@asynccontextmanager -async def _store_postgres_aio_pipe(): - if sys.version_info < (3, 10): - pytest.skip("Async Postgres tests require Python 3.10+") - database = f"test_{uuid4().hex[:16]}" - async with await AsyncConnection.connect( - DEFAULT_POSTGRES_URI, autocommit=True - ) as conn: - await conn.execute(f"CREATE DATABASE {database}") - try: - async with AsyncPostgresStore.from_conn_string( - DEFAULT_POSTGRES_URI + database - ) as store: - await store.setup() # Run in its own transaction - async with AsyncPostgresStore.from_conn_string( - DEFAULT_POSTGRES_URI + database, pipeline=True - ) as store: - yield store - finally: - async with await AsyncConnection.connect( - DEFAULT_POSTGRES_URI, autocommit=True - ) as conn: - await conn.execute(f"DROP DATABASE {database}") - - -@asynccontextmanager -async def _store_postgres_aio_pool(): - if sys.version_info < (3, 10): - pytest.skip("Async Postgres tests require Python 3.10+") - database = f"test_{uuid4().hex[:16]}" - async with await AsyncConnection.connect( - DEFAULT_POSTGRES_URI, autocommit=True - ) as conn: - await conn.execute(f"CREATE DATABASE {database}") - try: - async with AsyncPostgresStore.from_conn_string( - DEFAULT_POSTGRES_URI + database, - pool_config={"max_size": 10}, - ) as store: - await store.setup() - yield store - finally: - async with await AsyncConnection.connect( - DEFAULT_POSTGRES_URI, autocommit=True - ) as conn: - await conn.execute(f"DROP DATABASE {database}") - - @pytest.fixture(scope="function", params=["sqlite", "memory"]) def cache(request: pytest.FixtureRequest) -> Iterator[BaseCache]: if request.param == "sqlite": @@ -374,62 +311,6 @@ def cache(request: pytest.FixtureRequest) -> Iterator[BaseCache]: raise ValueError(f"Unknown cache type: {request.param}") -@contextmanager -def _store_postgres(): - database = f"test_{uuid4().hex[:16]}" - # create unique db - with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: - conn.execute(f"CREATE DATABASE {database}") - try: - # yield store - with PostgresStore.from_conn_string(DEFAULT_POSTGRES_URI + database) as store: - store.setup() - yield store - finally: - # drop unique db - with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: - conn.execute(f"DROP DATABASE {database}") - - -@contextmanager -def _store_postgres_pipe(): - database = f"test_{uuid4().hex[:16]}" - # create unique db - with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: - conn.execute(f"CREATE DATABASE {database}") - try: - # yield store - with PostgresStore.from_conn_string(DEFAULT_POSTGRES_URI + database) as store: - store.setup() # Run in its own transaction - with PostgresStore.from_conn_string( - DEFAULT_POSTGRES_URI + database, pipeline=True - ) as store: - yield store - finally: - # drop unique db - with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: - conn.execute(f"DROP DATABASE {database}") - - -@contextmanager -def _store_postgres_pool(): - database = f"test_{uuid4().hex[:16]}" - # create unique db - with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: - conn.execute(f"CREATE DATABASE {database}") - try: - # yield store - with PostgresStore.from_conn_string( - DEFAULT_POSTGRES_URI + database, pool_config={"max_size": 10} - ) as store: - store.setup() - yield store - finally: - # drop unique db - with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: - conn.execute(f"DROP DATABASE {database}") - - @pytest.fixture( scope="function", params=["in_memory", "postgres", "postgres_pipe", "postgres_pool"], @@ -439,7 +320,8 @@ def sync_store(request: pytest.FixtureRequest) -> Iterator[BaseStore]: if store_name is None: yield None elif store_name == "in_memory": - yield InMemoryStore() + with _store_memory() as store: + yield store elif store_name == "postgres": with _store_postgres() as store: yield store @@ -462,7 +344,8 @@ async def async_store(request: pytest.FixtureRequest) -> AsyncIterator[BaseStore if store_name is None: yield None elif store_name == "in_memory": - yield InMemoryStore() + with _store_memory() as store: + yield store elif store_name == "postgres_aio": async with _store_postgres_aio() as store: yield store diff --git a/libs/langgraph/tests/conftest_store.py b/libs/langgraph/tests/conftest_store.py new file mode 100644 index 000000000..9d047a8f8 --- /dev/null +++ b/libs/langgraph/tests/conftest_store.py @@ -0,0 +1,154 @@ +import sys +from contextlib import asynccontextmanager, contextmanager +from uuid import uuid4 + +import pytest +from psycopg import AsyncConnection, Connection + +from langgraph.store.memory import InMemoryStore +from langgraph.store.postgres import AsyncPostgresStore, PostgresStore + +DEFAULT_POSTGRES_URI = "postgres://postgres:postgres@localhost:5442/" + + +@contextmanager +def _store_memory(): + store = InMemoryStore() + yield store + + +@contextmanager +def _store_postgres(): + database = f"test_{uuid4().hex[:16]}" + # create unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"CREATE DATABASE {database}") + try: + # yield store + with PostgresStore.from_conn_string(DEFAULT_POSTGRES_URI + database) as store: + store.setup() + yield store + finally: + # drop unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"DROP DATABASE {database}") + + +@contextmanager +def _store_postgres_pipe(): + database = f"test_{uuid4().hex[:16]}" + # create unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"CREATE DATABASE {database}") + try: + # yield store + with PostgresStore.from_conn_string(DEFAULT_POSTGRES_URI + database) as store: + store.setup() # Run in its own transaction + with PostgresStore.from_conn_string( + DEFAULT_POSTGRES_URI + database, pipeline=True + ) as store: + yield store + finally: + # drop unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"DROP DATABASE {database}") + + +@contextmanager +def _store_postgres_pool(): + database = f"test_{uuid4().hex[:16]}" + # create unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"CREATE DATABASE {database}") + try: + # yield store + with PostgresStore.from_conn_string( + DEFAULT_POSTGRES_URI + database, pool_config={"max_size": 10} + ) as store: + store.setup() + yield store + finally: + # drop unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"DROP DATABASE {database}") + + +@asynccontextmanager +async def _store_postgres_aio(): + if sys.version_info < (3, 10): + pytest.skip("Async Postgres tests require Python 3.10+") + database = f"test_{uuid4().hex[:16]}" + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"CREATE DATABASE {database}") + try: + async with AsyncPostgresStore.from_conn_string( + DEFAULT_POSTGRES_URI + database + ) as store: + await store.setup() + yield store + finally: + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"DROP DATABASE {database}") + + +@asynccontextmanager +async def _store_postgres_aio_pipe(): + if sys.version_info < (3, 10): + pytest.skip("Async Postgres tests require Python 3.10+") + database = f"test_{uuid4().hex[:16]}" + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"CREATE DATABASE {database}") + try: + async with AsyncPostgresStore.from_conn_string( + DEFAULT_POSTGRES_URI + database + ) as store: + await store.setup() # Run in its own transaction + async with AsyncPostgresStore.from_conn_string( + DEFAULT_POSTGRES_URI + database, pipeline=True + ) as store: + yield store + finally: + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"DROP DATABASE {database}") + + +@asynccontextmanager +async def _store_postgres_aio_pool(): + if sys.version_info < (3, 10): + pytest.skip("Async Postgres tests require Python 3.10+") + database = f"test_{uuid4().hex[:16]}" + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"CREATE DATABASE {database}") + try: + async with AsyncPostgresStore.from_conn_string( + DEFAULT_POSTGRES_URI + database, + pool_config={"max_size": 10}, + ) as store: + await store.setup() + yield store + finally: + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"DROP DATABASE {database}") + + +__all__ = [ + "_store_memory", + "_store_postgres", + "_store_postgres_pipe", + "_store_postgres_pool", + "_store_postgres_aio", + "_store_postgres_aio_pipe", + "_store_postgres_aio_pool", +]