DEV Community

BAOFUFAN
BAOFUFAN

Posted on

Cut LLM Conversation Memory Regression Testing from 20 Minutes to 3 Seconds with Pytest + SQLite

Last Friday at 5 PM, just as I was about to shut down my machine and leave, a colleague dropped a message in the group chat: “The live conversation memory summary is missing again, the whole context is off.” My heart sank — I had just modified the session window truncation logic in that module the week before and shipped it after “roughly testing” it. Now that it was broken, I had to manually go back through the code scenario by scenario to fill in the gaps. By the time I finished, it was already dark, and I was left with a nagging feeling: “Did I really cover every branch this time?”

LLM conversation memory is not simple CRUD. It involves message sequencing, summary generation, vector retrieval, window truncation, and token budget management. Changing any one piece can silently break behavior that was previously stable. What makes it worse is memory bugs usually don’t crash outright — they silently drop information, and it’s hard to spot the problem in logs. Manual regression testing is incredibly expensive, and if you keep shipping by “gut feeling,” it’s only a matter of time before something blows up.

Later, I spent half a day building a complete suite of automated regression tests for the memory storage using Pytest + SQLite. Now 100+ test cases run in 3 seconds, and every code change is automatically guarded against silent message loss. This post pulls apart the whole solution and all the pitfalls I encountered, so you can drop it straight into your own project.

Breaking Down the Problem: Why Is Regression Testing Conversation Memory So Hard?

“Conversation memory” in LLM applications is rarely a simple message list. A typical memory module might include:

  1. Message storage: recording user / assistant / system messages by session.
  2. Context formatting: truncating history according to a token budget, keeping the most recent k turns.
  3. Summary generation: calling an LLM to generate a summary when the history gets too long, then injecting it into future prompts.
  4. Vector memory: generating embeddings for key information and using semantic retrieval to recall it (e.g., LangChain’s VectorStoreRetrieverMemory).

A logic change in any one of these parts can affect the final message list fed to the model. And the impact is hard to see with the naked eye — for example, you tweak the summary concatenation format, and in a certain type of conversation the summary doesn’t get inserted. If your testing doesn’t happen to trigger the summary condition, you’ll never catch it.

Previously, the team would tweak the code and then manually run a few typical conversations to inspect what the final assembled prompt looked like. The problems with that approach:

  • Branch coverage was up to chance. Changing part A could affect part B, but B never got tested.
  • Every time you had to manually construct a long conversation prefix, trigger a summary, and simulate token boundaries — very time-consuming.
  • For the vector retrieval side, manually constructing embedding results was completely unrealistic.

Why don’t the usual solutions work? If we used real LLMs and vector databases for testing, it would be slow and the results would be non-deterministic (LLMs are stochastic, embedding models can change when upgraded). Even worse, every test run consumes tokens, making the cost absurdly high. We needed a regression test suite that is fully deterministic, zero external dependencies, and executes in milliseconds.

Solution Design: Why Pytest + SQLite + Hand-rolled Stubs?

The core principle behind the technology choices: Everywhere the module under test has an external dependency, we abstract an interface and inject the dependency; in the test environment we replace it with a lightweight stub.

  • Pytest: goes without saying — its fixture system is a natural fit for managing memory instances used in tests.
  • SQLite: in-memory mode (:memory:) means zero configuration, zero residual state. Every test case gets a pristine database and runs extremely fast.
  • All LLM calls mocked: using unittest.mock to replace summary generation and embedding generation with deterministic fake functions that return preset results.
  • Vector retrieval replaced by a simple keyword-matching stub: it doesn’t touch a real vector DB; instead it does substring matching in-process, keeping behavior fully controllable.

Why didn’t we choose other solutions?

  • Not a real vector DB (like Chroma / Milvus): external services introduce network latency and non-determinism, turning tests from milliseconds into seconds while adding environment requirements.
  • Not integration tests instead of unit tests: integration tests can only tell you “is the final result correct?” — when something fails, it’s hard to tell whether the bug is in the memory logic or the LLM call. Isolating the business logic for pure logic testing is the real guarantee for regression safety.
  • Not file-based SQLite: :memory: gives the best performance and automatically cleans up after the test runs, leaving no garbage behind.

Architecturally, we split the memory module into two layers:

  1. MemoryStore: handles persistence and basic queries (SQLite implementation).
  2. MemoryManager: the business logic layer, responsible for assembling context, deciding when to summarize, and managing the token budget. It depends on MemoryStore and an LLMClient interface.

When testing, MemoryStore runs for real using SQLite in-memory, while LLMClient is completely mocked. This lets us verify the SQL logic in the storage layer while focusing on testing business decisions.

Core Implementation: Building a Usable Regression Test Suite Step by Step

First, define the core interfaces and the object under test. Suppose our MemoryManager looks something like this (simplified):

# memory.py
from dataclasses import dataclass
from typing import List, Optional

@dataclass
class Message:
    role: str        # "user" / "assistant" / "system"
    content: str

class MemoryStore:
    """持久化层,负责存和查"""
    def add_message(self, session_id: str, msg: Message): ...
    def get_history(self, session_id: str, limit: int = 10) -> List[Message]: ...
    def clear(self, session_id: str): ...

class LLMClient:
    """调用大模型(测试时会 mock)"""
    def generate_summary(self, messages: List[Message]) -> str: ...
    def get_embedding(self, text: str) -> List[float]: ...

class MemoryManager:
    def __init__(self, store: MemoryStore, llm: LLMClient,
                 max_tokens: int = 2000):
        self.store = store
        self.llm = llm
        self.max_tokens = max_tokens

    def build_context(self, session_id: str) -> List[Message]:
        # 核心逻辑:取最近消息 + 如果token超限则生成摘要
        ...
Enter fullscreen mode Exit fullscreen mode

First complete code block: SQLite implementation + Pytest fixture

Goal: Provide every test with a fully isolated and reusable memory store instance.

# conftest.py
import pytest
import sqlite3
from memory import MemoryStore, Message

class SQLiteMemoryStore(MemoryStore):
    def __init__(self, conn: sqlite3.Connection):
        self.conn = conn
        self.conn.execute(
            "CREATE TABLE IF NOT EXISTS messages "
            "(id INTEGER PRIMARY KEY AUTOINCREMENT, "
            "session_id TEXT, role TEXT, content TEXT, "
            "created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)"
        )

    def add_message(self, session_id: str, msg: Message):
        self.conn.execute(
            "INSERT INTO messages (session_id, role, content) VALUES (?, ?, ?)",
            (session_id, msg.role, msg.content)
        )
        self.conn.commit()

    def get_history(self, session_id: str, limit: int = 10) -> list[Message]:
        cur = self.conn.execute(
            "SELECT role, content FROM messages WHERE session_id=? "
            "ORDER BY created_at DESC LIMIT ?",
            (session_id, limit)
        )
        rows = cur.fetchall()
        return [Message(role=r[0], content=r[1]) for r in reversed(rows)]

    def clear(self, session_id: str):
        self.conn.execute(
            "DELETE FROM messages WHERE session_id=?", (session_id,)
        )
        self.conn.commit()

@pytest.fixture
def memory_store():
    conn = sqlite3.connect(":memory:")
    store = SQLiteMemoryStore(conn)
    yield store
    conn.close()
Enter fullscreen mode Exit fullscreen mode

Second complete code block: Mocking LLM calls — deterministic stubs

Goal: Replace summary generation and embeddings with deterministic functions so every test run produces the same result.

# test_helpers.py
from memory import LLMClient

def fake_summary(messages):
    """Return a fixed summary containing a hint about the input size."""
    return f"Summary of {len(messages)} messages."

def fake_embedding(text):
    """Return a deterministic pseudo-embedding; length matches real one."""
    # Simple approach: hash the text and produce a fixed-size list
    # In real use, return something like [0.1] * 1536
    return [0.1] * 3  # short for illustration

class StubLLMClient(LLMClient):
    def generate_summary(self, messages):
        return fake_summary(messages)

    def get_embedding(self, text):
        return fake_embedding(text)

@pytest.fixture
def stub_llm():
    return StubLLMClient()
Enter fullscreen mode Exit fullscreen mode

Third complete code block: Vector retrieval stub with keyword matching

Goal: Simulate vector retrieval that is fast, deterministic, and predictable.

# Suppose we have a VectorStore interface
class VectorStore:
    def add_memory(self, text: str, embedding: List[float]): ...
    def search(self, query: str, k: int = 3) -> List[str]: ...

class KeywordVectorStoreStub(VectorStore):
    def __init__(self):
        self.memories = []

    def add_memory(self, text: str, embedding=None):
        # Ignore embedding, just store text
        self.memories.append(text)

    def search(self, query: str, k: int = 3) -> List[str]:
        # Match by exact keyword containment (simple but deterministic)
        results = [m for m in self.memories if query.lower() in m.lower()]
        return results[:k]

@pytest.fixture
def vector_store_stub():
    return KeywordVectorStoreStub()
Enter fullscreen mode Exit fullscreen mode

Fourth complete code block: Test scenarios that actually caught bugs

Now we can write tests that cover tricky edge cases — the ones that used to silently pass incorrect contexts.

def test_summary_inserted_when_token_budget_exceeded(memory_store, stub_llm):
    manager = MemoryManager(memory_store, stub_llm, max_tokens=100)

    # Simulate a long conversation
    session = "test1"
    for i in range(20):
        memory_store.add_message(session, Message(
            role="user", content=f"Question {i} " + "x" * 20))
        memory_store.add_message(session, Message(
            role="assistant", content=f"Answer {i} " + "y" * 20))

    ctx = manager.build_context(session)
    # Expect a summary message at the beginning
    assert ctx[0].role == "system"
    assert "Summary of" in ctx[0].content
    # Recent messages should still be present
    assert ctx[-1].role == "assistant"

def test_no_summary_when_within_budget(memory_store, stub_llm):
    manager = MemoryManager(memory_store, stub_llm, max_tokens=10000)
    session = "test2"
    memory_store.add_message(session, Message(role="user", content="Hi"))
    memory_store.add_message(session, Message(role="assistant", content="Hello"))

    ctx = manager.build_context(session)
    # No summary, just the original messages
    assert len(ctx) == 2
    assert ctx[0].role == "user"

def test_old_message_truncated_before_summary(memory_store, stub_llm):
    manager = MemoryManager(memory_store, stub_llm, max_tokens=150)
    session = "test3"
    # Fill messages that definitely exceed budget
    for i in range(10):
        memory_store.add_message(session, Message(
            role="user", content=f"U{i} " + "z" * 30))
        memory_store.add_message(session, Message(
            role="assistant", content=f"A{i} " + "z" * 30))

    ctx = manager.build_context(session)
    # Early messages should not appear raw
    roles = [m.role for m in ctx]
    assert roles.count("system") == 1
    # The summary should reflect many messages, not just a couple
    assert "Summary of" in ctx[0].content

def test_vector_memory_retrieval_adds_relevant_context(memory_store, stub_llm, vector_store_stub):
    manager = MemoryManager(memory_store, stub_llm,
                            vector_store=vector_store_stub,
                            max_tokens=2000)
    # Pre-seed vector store with memories
    vector_store_stub.add_memory("User's name is Alice")
    vector_store_stub.add_memory("User lives in Berlin")

    session = "test4"
    memory_store.add_message(session, Message(
        role="user", content="What's my name?"))
    memory_store.add_message(session, Message(
        role="assistant", content="I don't know."))

    ctx = manager.build_context(session)
    # Should prepend the relevant memory
    # In a real manager, the vector retrieval result is injected
    # For illustration, we check if "Alice" appears in any system prompt
    system_msgs = [m.content for m in ctx if m.role == "system"]
    assert any("Alice" in s for s in system_msgs)
Enter fullscreen mode Exit fullscreen mode

With just these few tests, we immediately caught a regression: the summary concatenation format had changed, and the summary text was being mistakenly overwritten by the system prompt prefix.

Pitfalls and How to Steer Clear

While building this suite I ran into several non-obvious traps that could break the determinism guarantee or silently mask bugs.

1. Time-dependent logic in tests
The get_history method used CURRENT_TIMESTAMP, which is perfect for production but can cause order flakiness in tests that insert messages in rapid succession. The fix: explicitly control the insertion order, or use an auto-incrementing ID for ordering. In the stub, we used ORDER BY id DESC internally while the fixture uses created_at; but when created_at resolution is a second, simultaneous inserts get the same timestamp. Swap the ordering to rowid for tests, or insert with a tiny sleep, or better yet, refactor the store to accept a callable for the timestamp so tests can inject a controlled clock.

2. Not mocking the token counter
Token counting (len(encode(text))) introduces a hidden dependency on the tokenizer. If you use a real tokenizer, tests are sensitive to tokenizer library versions. Solution: mock the token counting function or use a simple character-length heuristic in tests. I injected a token_counter callable and replaced it with len in tests.

3. Leaking state between tests
SQLite :memory: is isolated per connection, but if the fixture accidentally shares the connection across tests you’ll get polluted state. The fixture above correctly creates a new connection per test. Double-check that your scope= is not set to "session" or "module" unless you know what you’re doing.

4. Over-mocking hides bugs
While we mock the LLM, we run the real SQLite store. This tests the storage layer thoroughly. If we mocked the store too, we’d lose coverage on SQL queries, filtering, and ordering. Keep the boundary at the right place: mock external I/O, keep internal state management real.

Impact: From 20 Minutes of Anxiety to 3 Seconds of Confidence

Before the test suite, a “quick check” meant manually running 5–6 scenarios, each requiring cURL or a test harness, then eyeballing the prompt. A full regression pass took at least 20 minutes and still missed paths.

After implementing the tests:

  • 100+ tests run in ~3 seconds.
  • Every commit triggers the full suite locally (and in CI).
  • The exact bug that prompted the Friday panic was caught in CI on the very next push — the diff that reintroduced the missing summary prefix got flagged immediately.
  • Developers now refactor memory logic with confidence, because the contracts are enforced by a fast, deterministic safety net.

Plug It Into Your Own Project

You can adapt this pattern to almost any LLM memory architecture:

  1. Identify the side effects — LLM calls, vector DB queries, time, filesystem.
  2. Define interfaces (Protocols or ABCs) for each external dependency.
  3. Write real implementation for production, and deterministic stubs for tests.
  4. Use Pytest fixtures to compose the test environment cleanly.
  5. Write characterization tests that capture current behavior before refactoring.

If you’re using LangChain, you can mock BaseChatModel and BaseRetriever in a similar way. For LlamaIndex, mock the LLM and VectorStoreIndex components.

The key takeaway: LLM applications are still software, and software needs fast, deterministic regression tests. Mock the non-deterministic parts, keep the business logic real, and you’ll never ship broken memory to production again.

Top comments (0)