From 30 Minutes to 5 Seconds: Testing Agent Memory with Pytest + SQLite

At 1 AM, a message popped up in the test group chat: the agent had leaked one user's memory into a new session. I opened DB Browser, hand-wrote 8 SELECT queries, and compared session_id and key row by row. Half an hour later, I tracked it down to a missing session_id in the WHERE condition of the upsert. This wasn't the first time.

Breaking Down the Problem

The agent's memory storage is basically just a SQLite table with three core fields: session_id, key, and value. The consistency requirements are simple too: updating the same key within a session must not create duplicate rows; different sessions must not bleed into each other; concurrent upserts must not overwrite each other.

But manual testing only covers whichever SQL query you happen to have in mind at the time. Each regression run requires preparing data, cleaning up old tables, and eyeballing the results—those three scenarios alone eat up 30 minutes. Even more frustrating, the occasional concurrent write issue is nearly impossible to reproduce by hand. The root cause is the lack of a fast, repeatable, automated consistency test—not that the business code itself is particularly complex.

Design

I chose pytest with an in-memory SQLite database. The core idea is that each test function gets a fresh :memory: store, and the connection closes automatically after the test.

I didn't use unittest because fixtures give you more flexible setup/teardown, and pytest's rewritten assertions produce clearer failure messages. I also skipped the production MySQL: it's slower and pulls in more environment dependencies. Since most of the consistency logic lives at the SQL layer, SQLite is enough to simulate it. The only extra thing needed is a file-backed fixture to verify that data survives a restart after disk persistence.

Core Implementation

This code implements a minimal but complete agent memory store. It uses a composite primary key (session_id, key) and an upsert with ON CONFLICT for updates. All the tests below build on it.

import sqlite3
import time
from contextlib import closing
from typing import Optional


class AgentMemoryStore:
    """Agent 记忆存储的最小实现，核心是复合主键和 upsert 一致性。"""

    def __init__(self, db_path: str = ":memory:", check_same_thread: bool = False):
        # check_same_thread=False 是为了后续并发测试不被线程检查打断
        self.conn = sqlite3.connect(db_path, check_same_thread=check_same_thread)
        self.conn.row_factory = sqlite3.Row
        self._init_schema()

    def _init_schema(self) -> None:
        with closing(self.conn.cursor()) as cur:
            cur.execute(
                """
                CREATE TABLE IF NOT EXISTS memories (
                    session_id TEXT NOT NULL,
                    key TEXT NOT NULL,
                    value TEXT NOT NULL,
                    updated_at REAL NOT NULL,
                    PRIMARY KEY (session_id, key)
                )
                """
            )
            self.conn.commit()

    def upsert(self, session_id: str, key: str, value: str) -> float:
        ts = time.time()
        with closing(self.conn.cursor()) as cur:
            cur.execute(
                """
                INSERT INTO memories (session_id, key, value, updated_at)
                VALUES (?, ?, ?, ?)
                ON CONFLICT(session_id, key)
                DO UPDATE SET value = excluded.value,
                              updated_at = excluded.updated_at
                """,
                (session_id, key, value, ts),
            )
            self.conn.commit()
        return ts

    def get(self, session_id: str, key: str) -> Optional[str]:
        with closing(self.conn.cursor()) as cur:
            cur.execute(
                "SELECT value FROM memories WHERE session_id = ? AND key = ?",
                (session_id, key),
            )
            row = cur.fetchone()
        return row["value"] if row else None

    def close(self) -> None:
        self.conn.close()

The test file below uses a function-scope fixture to give each test case its own database. It covers the three consistency points most likely to break: updating the same key, session isolation, and duplicate composite primary keys.

import pytest

from memory_store import AgentMemoryStore  # 假设上面的类放在 memory_store.py


@pytest.fixture()
def store():
    store = AgentMemoryStore(":memory:")
    yield store
    store.close()


def test_upsert_