More idiomatic sync/async store fixtures in pytest (#4617)

- Use a single fixture that returns the store, instead of list of names

libs/langgraph/tests/conftest.py

@@ -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,73 +311,41 @@ def cache(request: pytest.FixtureRequest) -> Iterator[BaseCache]:
         raise ValueError(f"Unknown cache type: {request.param}")
 
 
-@pytest.fixture(scope="function")
-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}")
-
-
-@pytest.fixture(scope="function")
-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}")
-
-
-@pytest.fixture(scope="function")
-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")
-def store_in_memory():
-    yield InMemoryStore()
-
-
-@asynccontextmanager
-async def awith_store(store_name: Optional[str]) -> AsyncIterator[BaseStore]:
+@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()
+        with _store_memory() as store:
+            yield store
+    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}")
+
+
+@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":
+        with _store_memory() as store:
+            yield store
     elif store_name == "postgres_aio":
         async with _store_postgres_aio() as store:
             yield store
@@ -483,15 +388,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",
-]

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",
+]

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()

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)