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

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

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)