Merge pull request #1579 from langchain-ai/nc/2sep/async-checkpoint-test-setup

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

libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py

@@ -1,6 +1,6 @@
 import asyncio
 from contextlib import asynccontextmanager
-from typing import Any, AsyncIterator, Optional, Union
+from typing import Any, AsyncIterator, Iterator, List, Optional, Union
 
 from langchain_core.runnables import RunnableConfig
 from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline
@@ -51,6 +51,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
         self.conn = conn
         self.pipe = pipe
         self.lock = asyncio.Lock()
+        self.loop = asyncio.get_running_loop()
 
     @classmethod
     @asynccontextmanager
@@ -329,3 +330,96 @@ class AsyncPostgresSaver(BasePostgresSaver):
                     binary=True, row_factory=dict_row
                 ) as cur:
                     yield cur
+
+    def list(
+        self,
+        config: Optional[RunnableConfig],
+        *,
+        filter: Optional[dict[str, Any]] = None,
+        before: Optional[RunnableConfig] = None,
+        limit: Optional[int] = None,
+    ) -> Iterator[CheckpointTuple]:
+        """List checkpoints from the database.
+
+        This method retrieves a list of checkpoint tuples from the Postgres database based
+        on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
+
+        Args:
+            config (Optional[RunnableConfig]): Base configuration for filtering checkpoints.
+            filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata.
+            before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None.
+            limit (Optional[int]): Maximum number of checkpoints to return.
+
+        Yields:
+            Iterator[CheckpointTuple]: An iterator of matching checkpoint tuples.
+        """
+        aiter_ = self.alist(config, filter=filter, before=before, limit=limit)
+        while True:
+            try:
+                yield asyncio.run_coroutine_threadsafe(
+                    anext(aiter_), self.loop
+                ).result()
+            except StopAsyncIteration:
+                break
+
+    def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
+        """Get a checkpoint tuple from the database.
+
+        This method retrieves a checkpoint tuple from the Postgres database based on the
+        provided config. If the config contains a "checkpoint_id" key, the checkpoint with
+        the matching thread ID and "checkpoint_id" is retrieved. Otherwise, the latest checkpoint
+        for the given thread ID is retrieved.
+
+        Args:
+            config (RunnableConfig): The config to use for retrieving the checkpoint.
+
+        Returns:
+            Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
+        """
+        return asyncio.run_coroutine_threadsafe(
+            self.aget_tuple(config), self.loop
+        ).result()
+
+    def put(
+        self,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+        new_versions: ChannelVersions,
+    ) -> RunnableConfig:
+        """Save a checkpoint to the database.
+
+        This method saves a checkpoint to the Postgres database. The checkpoint is associated
+        with the provided config and its parent config (if any).
+
+        Args:
+            config (RunnableConfig): The config to associate with the checkpoint.
+            checkpoint (Checkpoint): The checkpoint to save.
+            metadata (CheckpointMetadata): Additional metadata to save with the checkpoint.
+            new_versions (ChannelVersions): New channel versions as of this write.
+
+        Returns:
+            RunnableConfig: Updated configuration after storing the checkpoint.
+        """
+        return asyncio.run_coroutine_threadsafe(
+            self.aput(config, checkpoint, metadata, new_versions), self.loop
+        ).result()
+
+    def put_writes(
+        self,
+        config: RunnableConfig,
+        writes: List[tuple[str, Any]],
+        task_id: str,
+    ) -> None:
+        """Store intermediate writes linked to a checkpoint.
+
+        This method saves intermediate writes associated with a checkpoint to the database.
+
+        Args:
+            config (RunnableConfig): Configuration of the related checkpoint.
+            writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
+            task_id (str): Identifier for the task creating the writes.
+        """
+        return asyncio.run_coroutine_threadsafe(
+            self.aput_writes(config, writes, task_id), self.loop
+        ).result()

libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py

@@ -1,11 +1,11 @@
 import asyncio
-import functools
 from contextlib import asynccontextmanager
 from typing import (
     Any,
     AsyncIterator,
     Dict,
     Iterator,
+    List,
     Optional,
     Sequence,
     Tuple,
@@ -31,20 +31,6 @@ from langgraph.checkpoint.sqlite.utils import search_where
 T = TypeVar("T", bound=callable)
 
 
-def not_implemented_sync_method(func: T) -> T:
-    @functools.wraps(func)
-    def wrapper(*args, **kwargs):
-        raise NotImplementedError(
-            "The AsyncSqliteSaver does not support synchronous methods. "
-            "Consider using the SqliteSaver instead.\n"
-            "from langgraph.checkpoint.sqlite import SqliteSaver\n"
-            "See https://langchain-ai.github.io/langgraph/reference/checkpoints/langgraph.checkpoint.sqlite.SqliteSaver "
-            "for more information."
-        )
-
-    return wrapper
-
-
 class AsyncSqliteSaver(BaseCheckpointSaver):
     """An asynchronous checkpoint saver that stores checkpoints in a SQLite database.
 
@@ -132,6 +118,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
         self.jsonplus_serde = JsonPlusSerializer()
         self.conn = conn
         self.lock = asyncio.Lock()
+        self.loop = asyncio.get_running_loop()
         self.is_setup = False
 
     @classmethod
@@ -150,16 +137,24 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
         async with aiosqlite.connect(conn_string) as conn:
             yield AsyncSqliteSaver(conn)
 
-    @not_implemented_sync_method
     def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
         """Get a checkpoint tuple from the database.
 
-        Note:
-            This method is not implemented for the AsyncSqliteSaver. Use `aget` instead.
-            Or consider using the [SqliteSaver][sqlitesaver] checkpointer.
-        """
+        This method retrieves a checkpoint tuple from the SQLite database based on the
+        provided config. If the config contains a "checkpoint_id" key, the checkpoint with
+        the matching thread ID and checkpoint ID is retrieved. Otherwise, the latest checkpoint
+        for the given thread ID is retrieved.
+
+        Args:
+            config (RunnableConfig): The config to use for retrieving the checkpoint.
+
+        Returns:
+            Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
+        """
+        return asyncio.run_coroutine_threadsafe(
+            self.aget_tuple(config), self.loop
+        ).result()
 
-    @not_implemented_sync_method
     def list(
         self,
         config: Optional[RunnableConfig],
@@ -168,21 +163,60 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
         before: Optional[RunnableConfig] = None,
         limit: Optional[int] = None,
     ) -> Iterator[CheckpointTuple]:
-        """List checkpoints from the database.
+        """List checkpoints from the database asynchronously.
 
-        Note:
-            This method is not implemented for the AsyncSqliteSaver. Use `alist` instead.
-            Or consider using the [SqliteSaver][sqlitesaver] checkpointer.
+        This method retrieves a list of checkpoint tuples from the SQLite database based
+        on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
+
+        Args:
+            config (Optional[RunnableConfig]): Base configuration for filtering checkpoints.
+            filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata.
+            before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None.
+            limit (Optional[int]): Maximum number of checkpoints to return.
+
+        Yields:
+            Iterator[CheckpointTuple]: An iterator of matching checkpoint tuples.
         """
+        aiter_ = self.alist(config, filter=filter, before=before, limit=limit)
+        while True:
+            try:
+                yield asyncio.run_coroutine_threadsafe(
+                    anext(aiter_), self.loop
+                ).result()
+            except StopAsyncIteration:
+                break
 
-    @not_implemented_sync_method
     def put(
         self,
         config: RunnableConfig,
         checkpoint: Checkpoint,
         metadata: CheckpointMetadata,
+        new_versions: ChannelVersions,
     ) -> RunnableConfig:
-        """Save a checkpoint to the database. FOO"""
+        """Save a checkpoint to the database.
+
+        This method saves a checkpoint to the SQLite database. The checkpoint is associated
+        with the provided config and its parent config (if any).
+
+        Args:
+            config (RunnableConfig): The config to associate with the checkpoint.
+            checkpoint (Checkpoint): The checkpoint to save.
+            metadata (CheckpointMetadata): Additional metadata to save with the checkpoint.
+            new_versions (ChannelVersions): New channel versions as of this write.
+
+        Returns:
+            RunnableConfig: Updated configuration after storing the checkpoint.
+        """
+        return asyncio.run_coroutine_threadsafe(
+            self.aput(config, checkpoint, metadata, new_versions), self.loop
+        ).result()
+
+    def put_writes(
+        self, config: RunnableConfig, writes: List[Tuple[str, Any]], task_id: str
+    ) -> None:
+        return asyncio.run_coroutine_threadsafe(
+            self.aput_writes(config, writes, task_id), self.loop
+        ).result()
 
     async def setup(self) -> None:
         """Set up the checkpoint database asynchronously.

libs/langgraph/Makefile

@@ -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

libs/langgraph/poetry.lock

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

libs/langgraph/pyproject.toml

@@ -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
 #

libs/langgraph/tests/conftest.py

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

libs/langgraph/tests/test_channels.py

@@ -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:

libs/langgraph/tests/test_interruption.py

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

libs/langgraph/tests/test_prebuilt.py

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

libs/langgraph/tests/test_pregel.py

@@ -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

libs/langgraph/tests/test_state.py

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

libs/langgraph/tests/test_store.py

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

libs/langgraph/tests/test_utils.py

@@ -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: