Taking Over LLM Memory Store Testing with Pytest: 90% Fewer State Inconsistencies

At 2:17 a.m., a Slack message exploded with user feedback: “Your AI assistant acts like it has amnesia—it just stored meeting minutes, but when I ask again, it says nothing was saved.” I dug into the logs and found that the Redis write was acknowledged, but the entry in the vector store was empty. This “intermittent amnesia” had already pissed off the user for the third time. Previously, we’d just manually curl a few conversation turns and call that a test pass—clearly that approach was broken. I was forced to set a rule the next morning: All state‑consistency logic for memory storage must be automatically verified by Pytest, or it’s not allowed to land on main.

Problem Breakdown

An LLM’s “memory store” typically isn’t a single backend: short‑term memory lives in a Redis List or an in‑memory ring buffer, long‑term memory gets embedded and pushed to a vector database (Qdrant/Milvus), and there’s often a raw text copy in PostgreSQL. During a single save_memory call, if you just write to Redis and PG synchronously and then push to the vector store asynchronously, there’s a window where Redis has already returned but the vector store hasn’t persisted—and a read request during that window can get incomplete history.

Even nastier is “memory compaction.” When the conversation length exceeds a threshold, you trigger summary generation while also cleaning up original messages and updating summary vectors. If any step fails, you end up with both “compacted” and leftover original messages coexisting, leading to duplicate or contradictory content.

The usual remedy is adding a manual verify step or running reconciliation scripts in production to sweep up dirty data. This has three fatal flaws:

• Bugs during async windows are nearly impossible to reproduce manually; eyeballing logs is searching for a needle in a haystack.
• Reconciliation scripts are post‑mortem bandaids, not development‑stage safety nets.
• New team members have no idea which scenarios will break; the illusion of “it works” shatters at the slightest nudge.

So I decided to bake state‑consistency verification directly into a Pytest test suite and let CI do the heavy lifting.

Solution Design

In terms of tooling, the test framework was non‑negotiable: Pytest. The reasons are simple—fixtures elegantly manage the lifecycle of multiple storage backends, and parametrize can cover dozens of failure combinations in one go.

Why not the seemingly obvious unit tests?

• Unit tests mock a single layer and can’t verify cross‑storage concurrency races. Mocked return values always follow your script; they never reveal a real window.
• Integration tests with Docker Compose spin up real Redis and vector stores, but managing them directly in Pytest means writing a ton of boilerplate. We use pytest-docker or simply testcontainers-python so each test case gets isolated, fresh storage instances that are automatically destroyed after the run.

The architectural idea is straightforward: each test case follows a pattern of concurrent writes + multiple reads + eventual assertions. We simulate two coroutines simultaneously writing messages for the same user, then trigger a compaction signal, and finally poll all backends with an assert_eventually style until they are consistent. No guessing with time.sleep(3); we rely on tenacity’s retry or asyncio.wait_for for timed conditional waits.

I also explicitly modeled the memory store interface as a Protocol, so the test suite doesn’t just cover the current implementation—if we swap out the vector database or caching layer later, the exact same test cases can run with zero changes.

Core Implementation

1. Defining a Testable Memory Store Abstraction

This code answers the question “what interface do we actually test?” by abstracting the production‑level memory storage into a MemoryStoreProtocol. All subsequent tests depend only on this protocol.

from typing import Protocol, List, Dict, Any, runtime_checkable

@runtime_checkable
class MemoryStoreProtocol(Protocol):
    """大模型记忆存储必须实现的接口"""
    async def append_message(self, user_id: str, role: str, content: str) -> None:
        """追加一条对话消息到短期记忆"""
        ...

    async def get_history(self, user_id: str, limit: int = 10) -> List[Dict[str, Any]]:
        """按时间顺序返回最近的消息列表"""
        ...

    async def compact(self, user_id: str) -> None:
        """触发记忆压缩：生成摘要并清理原始短期消息"""
        ...

    async def consistent_state(self, user_id: str) -> bool:
        """检查所有后端数据是否一致（用于测试断言）"""
        ...

2. A Simplified Implementation with Race Conditions

To make the tests meaningful, I deliberately wrote a “buggy” implementation called FragileMemoryStore—it uses a Python list in place of Redis, a dict in place of the vector store, and randomly await asyncio.sleep(0) to simulate async delays. This code is the target we’re shooting at.

import asyncio
import random
from dataclasses import dataclass, field

@dataclass
class FragileMemoryStore:
    # 模拟三层存储
    short_term: Dict[str, List[dict]] = field(default_factory=dict)   # Redis List
    vector: Dict[str, List[str]] = field(default_factory=dict)       # 向量库
    raw_pg: Dict[str, List[dict]] = field(default_factory=dict)      # PG 原始表

    async def append_message(self, user_id: str, role: str, content: str) -> None:
        msg = {"role": role, "content": content}
        # 1. 先写短期缓存（同步写）
        self.short_term.setdefault(user_id, []).append(msg)
        # 2. 写 PG 原始表（同步写）
        self.raw_pg.setdefault(user_id, []).append(msg)
        # 3. 推向量库是异步的，模拟延迟
        await asyncio.sleep(random.uniform(0, 0.02))  # 随机窗口
        self.vector.setdefault(user_id,