From b7988c72df2deee870a564ccacf1d09955767131 Mon Sep 17 00:00:00 2001 From: 35C4n0r Date: Sun, 24 May 2026 10:04:56 +0530 Subject: [PATCH] test: add integration tests for all backends and cover known bug surfaces - test_policy_resolver: add two tests that exercise a policy_resolver creating fresh TTLPolicy instances per call; the 50-thread scale test reliably reproduces the id()-cache collision that caused wrong threads to be evaluated against the wrong TTL - test_sweeper: replace test_loop_fails_on_exception (which asserted the broken die-on-error behavior) with test_loop_continues_after_exception and test_loop_logs_exception that verify the loop survives transient failures and emits an ERROR-level log entry - test_strategy_sqlite: add tests that query the writes table directly after a sweep, and a cross-contamination test that verifies surviving threads are not touched - test_strategy_postgres: full rewrite - add _row_count helpers, collect tests, hard_age_ttl test, three-table cleanup assertions for both sync and async paths, advisory lock type checks, real two-sweeper lock contention test, and incremental cursor collection test - pyproject.toml: add test-postgres extra with testcontainers[postgres], langgraph-checkpoint-postgres, and psycopg[binary] - ci.yml: add test-postgres job that installs test-postgres extra and runs tests/test_strategy_postgres.py; exclude postgres tests from the matrix job so they only run under Docker --- .github/workflows/ci.yml | 25 ++- pyproject.toml | 5 + tests/test_policy_resolver.py | 50 ++++++ tests/test_strategy_postgres.py | 276 ++++++++++++++++++++++++++++---- tests/test_strategy_sqlite.py | 63 ++++++++ tests/test_sweeper.py | 42 +++-- 6 files changed, 423 insertions(+), 38 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a794671..ccc6602 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -32,4 +32,27 @@ jobs: run: uv sync --extra test - name: Run tests - run: uv run pytest + run: uv run pytest --ignore=tests/test_strategy_postgres.py + + test-postgres: + name: Postgres Integration Tests + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v6.0.2 + + - name: Install uv + uses: astral-sh/setup-uv@v8.1.0 + with: + enable-cache: true + + - name: Set up Python 3.12 + uses: actions/setup-python@v6.2.0 + with: + python-version: "3.12" + + - name: Install dependencies + run: uv sync --extra test --extra test-postgres + + - name: Run Postgres integration tests + run: uv run pytest tests/test_strategy_postgres.py -v diff --git a/pyproject.toml b/pyproject.toml index af40030..78fdeef 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,11 @@ test = [ "langgraph-checkpoint-sqlite>=3.1.0", "aiosqlite>=0.20", ] +test-postgres = [ + "langgraph-checkpoint-postgres>=3.1.0", + "psycopg[binary]>=3.1", + "testcontainers[postgres]>=4.0", +] benchmark = [ "langgraph-checkpoint-sqlite>=3.0.3", "langgraph-checkpoint-postgres>=3.0.5", diff --git a/tests/test_policy_resolver.py b/tests/test_policy_resolver.py index 26c31bd..0639d86 100644 --- a/tests/test_policy_resolver.py +++ b/tests/test_policy_resolver.py @@ -128,3 +128,53 @@ def resolver(tid): assert "exempt" not in result.deleted_thread_ids assert "expired" in result.deleted_thread_ids + + +def test_fresh_policy_objects_per_thread_applied_correctly(): + tuples = [ + make_checkpoint_tuple("strict", None, iso_ts(-70)), + make_checkpoint_tuple("lenient", None, iso_ts(-70)), + ] + cp = MockCP(tuples) + + def resolver(tid): + if tid == "strict": + return TTLPolicy(idle_ttl_seconds=60) + return TTLPolicy(idle_ttl_seconds=3600) + + sweeper = Sweeper( + cp, + TTLPolicy(idle_ttl_seconds=9999), + policy_resolver=resolver, + safe_delete=False, + _strategy=list_strategy(tuples), + ) + result = sweeper.sweep() + + assert "strict" in result.deleted_thread_ids + assert "lenient" not in result.deleted_thread_ids + + +def test_many_threads_fresh_policies_all_evaluated_correctly(): + thread_ids = [f"t{i}" for i in range(50)] + tuples = [make_checkpoint_tuple(tid, None, iso_ts(-70)) for tid in thread_ids] + cp = MockCP(tuples) + + strict_ids = {tid for tid in thread_ids if int(tid[1:]) % 2 == 0} + + def resolver(tid): + if tid in strict_ids: + return TTLPolicy(idle_ttl_seconds=60) + return TTLPolicy(idle_ttl_seconds=3600) + + sweeper = Sweeper( + cp, + TTLPolicy(idle_ttl_seconds=9999), + policy_resolver=resolver, + safe_delete=False, + _strategy=list_strategy(tuples), + ) + result = sweeper.sweep() + + deleted = set(result.deleted_thread_ids) + assert deleted == strict_ids diff --git a/tests/test_strategy_postgres.py b/tests/test_strategy_postgres.py index 429af00..22caeb3 100644 --- a/tests/test_strategy_postgres.py +++ b/tests/test_strategy_postgres.py @@ -1,13 +1,18 @@ """Integration tests for PostgresSaver / AsyncPostgresSaver strategies. -Requires Docker and testcontainers: - pip install testcontainers[postgres] +Requires the test-postgres extra: + uv sync --extra test --extra test-postgres """ import time +from typing import TypedDict import pytest from langgraph_ephemeral_checkpointer import Sweeper, TTLPolicy +from langgraph_ephemeral_checkpointer._coordination import ( + AsyncAdvisoryLock, + SyncAdvisoryLock, +) from langgraph_ephemeral_checkpointer._strategies import detect from langgraph_ephemeral_checkpointer._strategies.postgres import ( AsyncPostgresStrategy, @@ -16,8 +21,6 @@ testcontainers = pytest.importorskip("testcontainers") -from typing import TypedDict # noqa: E402 - from langgraph.graph import END, START, StateGraph # noqa: E402 from testcontainers.postgres import PostgresContainer # noqa: E402 # pyrefly: ignore[missing-import] @@ -25,6 +28,7 @@ class _State(TypedDict): x: int + def _build_graph(saver): builder = StateGraph(_State) # pyrefly: ignore[bad-specialization] builder.add_node("inc", lambda s: {"x": s["x"] + 1}) # pyrefly: ignore[bad-argument-type] @@ -32,27 +36,100 @@ def _build_graph(saver): builder.add_edge("inc", END) return builder.compile(checkpointer=saver) + +def _row_count(saver, table: str, thread_id: str) -> int: + with saver._cursor() as cur: + cur.execute(f"SELECT COUNT(*) AS c FROM {table} WHERE thread_id = %s", (thread_id,)) + return cur.fetchone()["c"] + + +async def _arow_count(saver, table: str, thread_id: str) -> int: + async with saver._cursor() as cur: + await cur.execute(f"SELECT COUNT(*) AS c FROM {table} WHERE thread_id = %s", (thread_id,)) + return (await cur.fetchone())["c"] + + @pytest.fixture(scope="module") def postgres_dsn(): with PostgresContainer("postgres:16-alpine") as pg: yield pg.get_connection_url().replace("psycopg2", "psycopg") + def test_detect_returns_postgres_strategy(postgres_dsn): from langgraph.checkpoint.postgres import PostgresSaver with PostgresSaver.from_conn_string(postgres_dsn) as saver: saver.setup() - strategy = detect(saver) - assert isinstance(strategy, PostgresStrategy) + assert isinstance(detect(saver), PostgresStrategy) + + +def test_detect_returns_async_postgres_strategy(postgres_dsn): + import asyncio + + from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver + + async def _inner(): + async with AsyncPostgresSaver.from_conn_string(postgres_dsn) as saver: + await saver.setup() + assert isinstance(detect(saver), AsyncPostgresStrategy) + + asyncio.run(_inner()) + def test_postgres_collect_empty(postgres_dsn): from langgraph.checkpoint.postgres import PostgresSaver with PostgresSaver.from_conn_string(postgres_dsn) as saver: saver.setup() - strategy = PostgresStrategy(saver) - threads, _ = strategy.collect(None) + threads, cursor = PostgresStrategy(saver).collect(None) assert isinstance(threads, dict) + assert cursor is None or isinstance(cursor, str) + + +def test_postgres_collect_returns_thread_timestamps(postgres_dsn): + from langgraph.checkpoint.postgres import PostgresSaver + + with PostgresSaver.from_conn_string(postgres_dsn) as saver: + saver.setup() + graph = _build_graph(saver) + graph.invoke({"x": 0}, {"configurable": {"thread_id": "collect_sync"}}) + + threads, cursor = PostgresStrategy(saver).collect(None) + + assert "collect_sync" in threads + ts = threads["collect_sync"] + assert ts.earliest_id <= ts.latest_id + assert cursor == ts.latest_id + + +@pytest.mark.asyncio +async def test_async_postgres_collect_returns_thread_timestamps(postgres_dsn): + from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver + + async with AsyncPostgresSaver.from_conn_string(postgres_dsn) as saver: + await saver.setup() + graph = _build_graph(saver) + await graph.ainvoke({"x": 0}, {"configurable": {"thread_id": "collect_async"}}) + + threads, cursor = await AsyncPostgresStrategy(saver).acollect(None) + + assert "collect_async" in threads + ts = threads["collect_async"] + assert ts.earliest_id <= ts.latest_id + assert cursor == ts.latest_id + + +def test_postgres_sweep_keeps_active_thread(postgres_dsn): + from langgraph.checkpoint.postgres import PostgresSaver + + with PostgresSaver.from_conn_string(postgres_dsn) as saver: + saver.setup() + graph = _build_graph(saver) + graph.invoke({"x": 0}, {"configurable": {"thread_id": "active_sync"}}) + + result = Sweeper(saver, TTLPolicy(idle_ttl_seconds=60)).sweep() + assert "active_sync" not in result.deleted_thread_ids + def test_postgres_sweep_deletes_expired(postgres_dsn): from langgraph.checkpoint.postgres import PostgresSaver @@ -60,43 +137,186 @@ def test_postgres_sweep_deletes_expired(postgres_dsn): with PostgresSaver.from_conn_string(postgres_dsn) as saver: saver.setup() graph = _build_graph(saver) - graph.invoke({"x": 0}, {"configurable": {"thread_id": "t1"}}) + graph.invoke({"x": 0}, {"configurable": {"thread_id": "expired_sync"}}) sweeper = Sweeper(saver, TTLPolicy(idle_ttl_seconds=1)) - - result = sweeper.sweep() - assert "t1" not in result.deleted_thread_ids + assert "expired_sync" not in sweeper.sweep().deleted_thread_ids time.sleep(1.1) - result = sweeper.sweep() - assert "t1" in result.deleted_thread_ids + assert "expired_sync" in sweeper.sweep().deleted_thread_ids -def test_detect_returns_async_postgres_strategy(postgres_dsn): + +@pytest.mark.asyncio +async def test_async_postgres_sweep_deletes_expired(postgres_dsn): from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver - async def _inner(): - async with AsyncPostgresSaver.from_conn_string(postgres_dsn) as saver: - await saver.setup() - strategy = detect(saver) - assert isinstance(strategy, AsyncPostgresStrategy) + async with AsyncPostgresSaver.from_conn_string(postgres_dsn) as saver: + await saver.setup() + graph = _build_graph(saver) + await graph.ainvoke({"x": 0}, {"configurable": {"thread_id": "expired_async"}}) + + sweeper = Sweeper(saver, TTLPolicy(idle_ttl_seconds=1)) + assert "expired_async" not in (await sweeper.asweep()).deleted_thread_ids + + time.sleep(1.1) + assert "expired_async" in (await sweeper.asweep()).deleted_thread_ids + + +def test_postgres_batch_delete_clears_all_tables(postgres_dsn): + from langgraph.checkpoint.postgres import PostgresSaver + + with PostgresSaver.from_conn_string(postgres_dsn) as saver: + saver.setup() + graph = _build_graph(saver) + graph.invoke({"x": 0}, {"configurable": {"thread_id": "all_tables_sync"}}) + + assert _row_count(saver, "checkpoints", "all_tables_sync") > 0 + assert _row_count(saver, "checkpoint_blobs", "all_tables_sync") > 0 + + sweeper = Sweeper(saver, TTLPolicy(idle_ttl_seconds=1)) + time.sleep(1.1) + result = sweeper.sweep() + + assert "all_tables_sync" in result.deleted_thread_ids + assert _row_count(saver, "checkpoints", "all_tables_sync") == 0 + assert _row_count(saver, "checkpoint_writes", "all_tables_sync") == 0 + assert _row_count(saver, "checkpoint_blobs", "all_tables_sync") == 0 - import asyncio - asyncio.run(_inner()) @pytest.mark.asyncio -async def test_async_postgres_sweep(postgres_dsn): +async def test_async_postgres_batch_delete_clears_all_tables(postgres_dsn): from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver async with AsyncPostgresSaver.from_conn_string(postgres_dsn) as saver: await saver.setup() graph = _build_graph(saver) - await graph.ainvoke({"x": 0}, {"configurable": {"thread_id": "t1"}}) + await graph.ainvoke({"x": 0}, {"configurable": {"thread_id": "all_tables_async"}}) - sweeper = Sweeper(saver, TTLPolicy(idle_ttl_seconds=1)) + assert await _arow_count(saver, "checkpoints", "all_tables_async") > 0 + assert await _arow_count(saver, "checkpoint_blobs", "all_tables_async") > 0 + sweeper = Sweeper(saver, TTLPolicy(idle_ttl_seconds=1)) + time.sleep(1.1) result = await sweeper.asweep() - assert "t1" not in result.deleted_thread_ids + + assert "all_tables_async" in result.deleted_thread_ids + assert await _arow_count(saver, "checkpoints", "all_tables_async") == 0 + assert await _arow_count(saver, "checkpoint_writes", "all_tables_async") == 0 + assert await _arow_count(saver, "checkpoint_blobs", "all_tables_async") == 0 + + +def test_postgres_sweep_no_cross_contamination(postgres_dsn): + from langgraph.checkpoint.postgres import PostgresSaver + + with PostgresSaver.from_conn_string(postgres_dsn) as saver: + saver.setup() + graph = _build_graph(saver) + graph.invoke({"x": 0}, {"configurable": {"thread_id": "survivor"}}) + graph.invoke({"x": 0}, {"configurable": {"thread_id": "victim"}}) time.sleep(1.1) - result = await sweeper.asweep() - assert "t1" in result.deleted_thread_ids + graph.invoke({"x": 1}, {"configurable": {"thread_id": "survivor"}}) + + result = Sweeper(saver, TTLPolicy(idle_ttl_seconds=1)).sweep() + + assert "victim" in result.deleted_thread_ids + assert "survivor" not in result.deleted_thread_ids + assert _row_count(saver, "checkpoints", "survivor") > 0 + assert _row_count(saver, "checkpoint_blobs", "survivor") > 0 + + +def test_postgres_hard_age_ttl(postgres_dsn): + from langgraph.checkpoint.postgres import PostgresSaver + + with PostgresSaver.from_conn_string(postgres_dsn) as saver: + saver.setup() + graph = _build_graph(saver) + graph.invoke({"x": 0}, {"configurable": {"thread_id": "hard_age"}}) + + sweeper = Sweeper(saver, TTLPolicy(hard_age_ttl_seconds=1)) + assert "hard_age" not in sweeper.sweep().deleted_thread_ids + + time.sleep(1.1) + assert "hard_age" in sweeper.sweep().deleted_thread_ids + + +def test_advisory_lock_type_for_postgres(postgres_dsn): + from langgraph.checkpoint.postgres import PostgresSaver + + with PostgresSaver.from_conn_string(postgres_dsn) as saver: + saver.setup() + sweeper = Sweeper(saver, TTLPolicy(idle_ttl_seconds=60), enable_coordination=True) + assert isinstance(sweeper._advisory_lock, SyncAdvisoryLock) + + +@pytest.mark.asyncio +async def test_advisory_lock_type_for_async_postgres(postgres_dsn): + from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver + + async with AsyncPostgresSaver.from_conn_string(postgres_dsn) as saver: + await saver.setup() + sweeper = Sweeper(saver, TTLPolicy(idle_ttl_seconds=60), enable_coordination=True) + assert isinstance(sweeper._advisory_lock, AsyncAdvisoryLock) + + +def test_postgres_advisory_lock_second_sweeper_skips(postgres_dsn): + from langgraph.checkpoint.postgres import PostgresSaver + + with PostgresSaver.from_conn_string(postgres_dsn) as s1: + with PostgresSaver.from_conn_string(postgres_dsn) as s2: + s1.setup() + + sw1 = Sweeper(s1, TTLPolicy(idle_ttl_seconds=60), enable_coordination=True) + sw2 = Sweeper(s2, TTLPolicy(idle_ttl_seconds=60), enable_coordination=True) + + assert sw1._advisory_lock.try_acquire() + try: + result = sw2.sweep() + assert result.deleted_thread_ids == [] + assert result.active_thread_count == 0 + finally: + sw1._advisory_lock.release() + + assert sw2.sweep() is not None + + +@pytest.mark.asyncio +async def test_async_postgres_advisory_lock_second_sweeper_skips(postgres_dsn): + from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver + + async with AsyncPostgresSaver.from_conn_string(postgres_dsn) as s1: + async with AsyncPostgresSaver.from_conn_string(postgres_dsn) as s2: + await s1.setup() + + sw1 = Sweeper(s1, TTLPolicy(idle_ttl_seconds=60), enable_coordination=True) + sw2 = Sweeper(s2, TTLPolicy(idle_ttl_seconds=60), enable_coordination=True) + + assert await sw1._advisory_lock.atry_acquire() + try: + result = await sw2.asweep() + assert result.deleted_thread_ids == [] + assert result.active_thread_count == 0 + finally: + await sw1._advisory_lock.arelease() + + +def test_postgres_incremental_collect(postgres_dsn): + from langgraph.checkpoint.postgres import PostgresSaver + + with PostgresSaver.from_conn_string(postgres_dsn) as saver: + saver.setup() + graph = _build_graph(saver) + graph.invoke({"x": 0}, {"configurable": {"thread_id": "cursor_t1"}}) + + strategy = PostgresStrategy(saver) + threads1, cursor1 = strategy.collect(None) + assert "cursor_t1" in threads1 + assert cursor1 is not None + + threads2, _ = strategy.collect(cursor1) + assert "cursor_t1" not in threads2 + + graph.invoke({"x": 0}, {"configurable": {"thread_id": "cursor_t2"}}) + threads3, _ = strategy.collect(cursor1) + assert "cursor_t2" in threads3 + assert "cursor_t1" not in threads3 diff --git a/tests/test_strategy_sqlite.py b/tests/test_strategy_sqlite.py index 4e19168..61079c6 100644 --- a/tests/test_strategy_sqlite.py +++ b/tests/test_strategy_sqlite.py @@ -112,3 +112,66 @@ async def test_async_sweep_deletes_expired(): time.sleep(1.1) result = await sweeper.asweep() assert "t1" in result.deleted_thread_ids + + +def test_sqlite_batch_delete_clears_writes_table(): + with SqliteSaver.from_conn_string(":memory:") as saver: + graph = _build_graph(saver) + graph.invoke({"x": 0}, {"configurable": {"thread_id": "t_writes"}}) + graph.invoke({"x": 1}, {"configurable": {"thread_id": "t_writes"}}) + + sweeper = Sweeper(saver, TTLPolicy(idle_ttl_seconds=1)) + time.sleep(1.1) + result = sweeper.sweep() + + assert "t_writes" in result.deleted_thread_ids + assert saver.conn.execute( + "SELECT COUNT(*) FROM checkpoints WHERE thread_id = ?", ("t_writes",) + ).fetchone()[0] == 0 + assert saver.conn.execute( + "SELECT COUNT(*) FROM writes WHERE thread_id = ?", ("t_writes",) + ).fetchone()[0] == 0 + + +@pytest.mark.asyncio +async def test_async_sqlite_batch_delete_clears_writes_table(): + pytest.importorskip("aiosqlite") + + async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: + graph = _build_graph(saver) + await graph.ainvoke({"x": 0}, {"configurable": {"thread_id": "t_writes_async"}}) + + sweeper = Sweeper(saver, TTLPolicy(idle_ttl_seconds=1)) + time.sleep(1.1) + result = await sweeper.asweep() + + assert "t_writes_async" in result.deleted_thread_ids + + async with saver.conn.execute( + "SELECT COUNT(*) FROM checkpoints WHERE thread_id = ?", ("t_writes_async",) + ) as cur: + assert (await cur.fetchone())[0] == 0 + + async with saver.conn.execute( + "SELECT COUNT(*) FROM writes WHERE thread_id = ?", ("t_writes_async",) + ) as cur: + assert (await cur.fetchone())[0] == 0 + + +def test_sqlite_sweep_no_cross_contamination(): + with SqliteSaver.from_conn_string(":memory:") as saver: + graph = _build_graph(saver) + graph.invoke({"x": 0}, {"configurable": {"thread_id": "old"}}) + graph.invoke({"x": 0}, {"configurable": {"thread_id": "new"}}) + + time.sleep(1.1) + graph.invoke({"x": 1}, {"configurable": {"thread_id": "new"}}) + + sweeper = Sweeper(saver, TTLPolicy(idle_ttl_seconds=1)) + result = sweeper.sweep() + + assert "old" in result.deleted_thread_ids + assert "new" not in result.deleted_thread_ids + assert saver.conn.execute( + "SELECT COUNT(*) FROM checkpoints WHERE thread_id = ?", ("new",) + ).fetchone()[0] > 0 diff --git a/tests/test_sweeper.py b/tests/test_sweeper.py index 06f90a8..08e540a 100644 --- a/tests/test_sweeper.py +++ b/tests/test_sweeper.py @@ -161,22 +161,46 @@ async def test_start_twice_raises(): await sweeper.stop() @pytest.mark.asyncio -async def test_loop_fails_on_exception(): - """An exception in asweep() fails the background task.""" +async def test_loop_continues_after_exception(): + call_count = 0 + + async def maybe_fail(self, *, dry_run=False): + nonlocal call_count + call_count += 1 + if call_count == 1: + raise RuntimeError("transient failure") + return SweepResult(deleted_thread_ids=[], active_thread_count=0, sweep_duration_seconds=0.0) + + cp = MockCheckpointer([]) + sweeper = Sweeper(cp, TTLPolicy(idle_ttl_seconds=60), _strategy=list_strategy([])) + + with patch.object(Sweeper, "asweep", maybe_fail): + await sweeper.start(interval_seconds=0) + await asyncio.sleep(0.1) + + assert sweeper._task is not None + assert not sweeper._task.done() + assert call_count >= 2 + await sweeper.stop() + + +@pytest.mark.asyncio +async def test_loop_logs_exception(caplog): + import logging + async def failing_asweep(self, *, dry_run=False): - raise RuntimeError("sweep failed") + raise RuntimeError("database timeout") cp = MockCheckpointer([]) sweeper = Sweeper(cp, TTLPolicy(idle_ttl_seconds=60), _strategy=list_strategy([])) with patch.object(Sweeper, "asweep", failing_asweep): - await sweeper.start(interval_seconds=600) - await asyncio.sleep(0.05) + with caplog.at_level(logging.ERROR, logger="langgraph_ephemeral_checkpointer.sweeper"): + await sweeper.start(interval_seconds=600) + await asyncio.sleep(0.05) + await sweeper.stop() - assert sweeper._task is not None - assert sweeper._task.done() - with pytest.raises(RuntimeError, match="sweep failed"): - sweeper._task.result() + assert any("Sweep cycle failed" in r.message for r in caplog.records) @pytest.mark.asyncio async def test_stop_before_start_is_noop():