"""Tests for shallow SQLite checkpoint savers (single-checkpoint-per-thread mode). Uses in-memory SQLite (``:memory:``) — no filesystem dependency. """ from __future__ import annotations import pytest # --------------------------------------------------------------------------- # AsyncShallowSqliteSaver tests # --------------------------------------------------------------------------- class TestAsyncShallowSqliteSaver: """Tests for ``AsyncShallowSqliteSaver`` — async shallow persistence.""" @pytest.mark.anyio async def test_aput_deletes_old_checkpoints_before_insert(self): """After two aput calls for the same thread, only 1 checkpoint remains.""" from deerflow.agents.checkpointer.shallow_sqlite import _make_async_shallow_saver ShallowSaver = _make_async_shallow_saver() async with ShallowSaver.from_conn_string(":memory:") as saver: await saver.setup() thread_config = {"configurable": {"thread_id": "test-thread-1", "checkpoint_ns": ""}} checkpoint_1 = {"ts": "2024-01-01T00:00:00Z", "id": "ckpt-1", "channel_values": {"x": 1}} checkpoint_2 = {"ts": "2024-01-01T00:01:00Z", "id": "ckpt-2", "channel_values": {"x": 2}} # Write first checkpoint await saver.aput(thread_config, checkpoint_1, {"source": "input", "step": 1, "writes": {}}, {}) # Write second checkpoint — should delete the first await saver.aput(thread_config, checkpoint_2, {"source": "loop", "step": 2, "writes": {}}, {}) # Verify only 1 checkpoint remains results = [] async for ckpt in saver.alist(thread_config): results.append(ckpt) assert len(results) == 1, f"Expected 1 checkpoint, got {len(results)}" assert results[0].config["configurable"]["checkpoint_id"] == "ckpt-2" @pytest.mark.anyio async def test_different_threads_do_not_interfere(self): """Checkpoints for different thread_ids are independent.""" from deerflow.agents.checkpointer.shallow_sqlite import _make_async_shallow_saver ShallowSaver = _make_async_shallow_saver() async with ShallowSaver.from_conn_string(":memory:") as saver: await saver.setup() t1 = {"configurable": {"thread_id": "thread-A", "checkpoint_ns": ""}} t2 = {"configurable": {"thread_id": "thread-B", "checkpoint_ns": ""}} await saver.aput(t1, {"ts": "Z", "id": "a1", "channel_values": {}}, {"source": "input", "step": 1, "writes": {}}, {}) await saver.aput(t2, {"ts": "Z", "id": "b1", "channel_values": {}}, {"source": "input", "step": 1, "writes": {}}, {}) await saver.aput(t1, {"ts": "Z", "id": "a2", "channel_values": {}}, {"source": "loop", "step": 2, "writes": {}}, {}) # Thread A: only ckpt a2 a_results = [c async for c in saver.alist(t1)] assert len(a_results) == 1 assert a_results[0].config["configurable"]["checkpoint_id"] == "a2" # Thread B: still has b1 b_results = [c async for c in saver.alist(t2)] assert len(b_results) == 1 assert b_results[0].config["configurable"]["checkpoint_id"] == "b1" @pytest.mark.anyio async def test_writes_table_also_cleaned(self): """aput_writes entries from old checkpoints are also deleted.""" from deerflow.agents.checkpointer.shallow_sqlite import _make_async_shallow_saver ShallowSaver = _make_async_shallow_saver() async with ShallowSaver.from_conn_string(":memory:") as saver: await saver.setup() thread_config = {"configurable": {"thread_id": "test-writes", "checkpoint_ns": ""}} # Write checkpoint 1 with associated writes ckpt1_config = await saver.aput( thread_config, {"ts": "Z", "id": "ckpt-w1", "channel_values": {}}, {"source": "input", "step": 1, "writes": {}}, {}, ) await saver.aput_writes(ckpt1_config, [("messages", "hello")], "task-1", "") # Write checkpoint 2 — should delete ckpt1 + its writes ckpt2_config = await saver.aput( thread_config, {"ts": "Z", "id": "ckpt-w2", "channel_values": {}}, {"source": "loop", "step": 2, "writes": {}}, {}, ) await saver.aput_writes(ckpt2_config, [("messages", "world")], "task-2", "") # Verify only 1 checkpoint remains results = [c async for c in saver.alist(thread_config)] assert len(results) == 1 assert results[0].config["configurable"]["checkpoint_id"] == "ckpt-w2" # Verify only the latest writes exist (get_tuple returns writes) latest = await saver.aget_tuple(ckpt2_config) assert latest is not None @pytest.mark.anyio async def test_full_mode_retains_all_checkpoints(self): """In full mode (default), all checkpoints are preserved.""" from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: await saver.setup() thread_config = {"configurable": {"thread_id": "test-full", "checkpoint_ns": ""}} await saver.aput(thread_config, {"ts": "Z", "id": "f1", "channel_values": {}}, {"source": "input", "step": 1, "writes": {}}, {}) await saver.aput(thread_config, {"ts": "Z", "id": "f2", "channel_values": {}}, {"source": "loop", "step": 2, "writes": {}}, {}) results = [c async for c in saver.alist(thread_config)] assert len(results) == 2, "Full mode should retain all checkpoints" # --------------------------------------------------------------------------- # ShallowSqliteSaver (sync) tests # --------------------------------------------------------------------------- class TestShallowSqliteSaver: """Tests for ``ShallowSqliteSaver`` — sync shallow persistence.""" def test_put_deletes_old_checkpoints_before_insert(self): """After two put calls for the same thread, only 1 checkpoint remains.""" from deerflow.agents.checkpointer.shallow_sqlite import _make_sync_shallow_saver ShallowSaver = _make_sync_shallow_saver() with ShallowSaver.from_conn_string(":memory:") as saver: saver.setup() thread_config = {"configurable": {"thread_id": "test-sync-1", "checkpoint_ns": ""}} checkpoint_1 = {"ts": "2024-01-01T00:00:00Z", "id": "ckpt-s1", "channel_values": {"x": 1}} checkpoint_2 = {"ts": "2024-01-01T00:01:00Z", "id": "ckpt-s2", "channel_values": {"x": 2}} saver.put(thread_config, checkpoint_1, {"source": "input", "step": 1, "writes": {}}, {}) saver.put(thread_config, checkpoint_2, {"source": "loop", "step": 2, "writes": {}}, {}) results = list(saver.list(thread_config)) assert len(results) == 1, f"Expected 1 checkpoint, got {len(results)}" assert results[0].config["configurable"]["checkpoint_id"] == "ckpt-s2" def test_different_threads_do_not_interfere_sync(self): """Checkpoints for different thread_ids are independent (sync).""" from deerflow.agents.checkpointer.shallow_sqlite import _make_sync_shallow_saver ShallowSaver = _make_sync_shallow_saver() with ShallowSaver.from_conn_string(":memory:") as saver: saver.setup() t1 = {"configurable": {"thread_id": "thread-A", "checkpoint_ns": ""}} t2 = {"configurable": {"thread_id": "thread-B", "checkpoint_ns": ""}} saver.put(t1, {"ts": "Z", "id": "a1", "channel_values": {}}, {"source": "input", "step": 1, "writes": {}}, {}) saver.put(t2, {"ts": "Z", "id": "b1", "channel_values": {}}, {"source": "input", "step": 1, "writes": {}}, {}) saver.put(t1, {"ts": "Z", "id": "a2", "channel_values": {}}, {"source": "loop", "step": 2, "writes": {}}, {}) a_results = list(saver.list(t1)) assert len(a_results) == 1 assert a_results[0].config["configurable"]["checkpoint_id"] == "a2" b_results = list(saver.list(t2)) assert len(b_results) == 1 assert b_results[0].config["configurable"]["checkpoint_id"] == "b1" def test_writes_table_also_cleaned_sync(self): """put_writes entries from old checkpoints are also deleted (sync).""" from deerflow.agents.checkpointer.shallow_sqlite import _make_sync_shallow_saver ShallowSaver = _make_sync_shallow_saver() with ShallowSaver.from_conn_string(":memory:") as saver: saver.setup() thread_config = {"configurable": {"thread_id": "test-sync-writes", "checkpoint_ns": ""}} ckpt1_config = saver.put( thread_config, {"ts": "Z", "id": "ckpt-sw1", "channel_values": {}}, {"source": "input", "step": 1, "writes": {}}, {}, ) saver.put_writes(ckpt1_config, [("messages", "hello")], "task-1", "") ckpt2_config = saver.put( thread_config, {"ts": "Z", "id": "ckpt-sw2", "channel_values": {}}, {"source": "loop", "step": 2, "writes": {}}, {}, ) saver.put_writes(ckpt2_config, [("messages", "world")], "task-2", "") results = list(saver.list(thread_config)) assert len(results) == 1 assert results[0].config["configurable"]["checkpoint_id"] == "ckpt-sw2" # --------------------------------------------------------------------------- # Config integration tests # --------------------------------------------------------------------------- class TestShallowConfig: """Tests for configuration and factory integration.""" def test_sqlite_mode_defaults_to_full(self): """sqlite_mode defaults to 'full' when not specified.""" from deerflow.config.checkpointer_config import CheckpointerConfig config = CheckpointerConfig(type="sqlite", connection_string="test.db") assert config.sqlite_mode == "full" def test_sqlite_mode_shallow_accepted(self): """sqlite_mode can be set to 'shallow'.""" from deerflow.config.checkpointer_config import CheckpointerConfig config = CheckpointerConfig(type="sqlite", connection_string="test.db", sqlite_mode="shallow") assert config.sqlite_mode == "shallow" def test_load_sqlite_config_with_shallow_mode(self): """load_checkpointer_config_from_dict accepts sqlite_mode.""" from deerflow.config.checkpointer_config import ( get_checkpointer_config, load_checkpointer_config_from_dict, set_checkpointer_config, ) set_checkpointer_config(None) load_checkpointer_config_from_dict({ "type": "sqlite", "connection_string": "/tmp/test.db", "sqlite_mode": "shallow", }) config = get_checkpointer_config() assert config is not None assert config.sqlite_mode == "shallow" def test_load_sqlite_config_defaults_sqlite_mode(self): """load_checkpointer_config_from_dict defaults sqlite_mode to 'full' when omitted.""" from deerflow.config.checkpointer_config import ( get_checkpointer_config, load_checkpointer_config_from_dict, set_checkpointer_config, ) set_checkpointer_config(None) load_checkpointer_config_from_dict({ "type": "sqlite", "connection_string": "/tmp/test.db", }) config = get_checkpointer_config() assert config is not None assert config.sqlite_mode == "full" @pytest.mark.anyio async def test_async_factory_uses_shallow_saver(self): """When sqlite_mode=shallow, async factory returns AsyncShallowSqliteSaver.""" from deerflow.agents.checkpointer.async_provider import _async_checkpointer from deerflow.config.checkpointer_config import CheckpointerConfig, set_checkpointer_config set_checkpointer_config(CheckpointerConfig( type="sqlite", connection_string=":memory:", sqlite_mode="shallow", )) config = CheckpointerConfig(type="sqlite", connection_string=":memory:", sqlite_mode="shallow") async with _async_checkpointer(config) as saver: # Should be an instance of the shallow saver cls_name = type(saver).__name__ assert "Shallow" in cls_name, f"Expected shallow saver, got {cls_name}" set_checkpointer_config(None) def test_sync_factory_uses_shallow_saver(self): """When sqlite_mode=shallow, sync factory returns ShallowSqliteSaver.""" from deerflow.agents.checkpointer.provider import _sync_checkpointer_cm from deerflow.config.checkpointer_config import CheckpointerConfig config = CheckpointerConfig(type="sqlite", connection_string=":memory:", sqlite_mode="shallow") with _sync_checkpointer_cm(config) as saver: cls_name = type(saver).__name__ assert "Shallow" in cls_name, f"Expected shallow saver, got {cls_name}" def test_invalid_sqlite_mode_raises(self): """Invalid sqlite_mode value raises validation error.""" from deerflow.config.checkpointer_config import CheckpointerConfig with pytest.raises(Exception): CheckpointerConfig(type="sqlite", connection_string="test.db", sqlite_mode="unknown")