From dd55e5097b06c6c7a729b8a7510a6a5809b146d8 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 9 May 2025 09:01:25 -0700 Subject: [PATCH] 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", +]