How I Squashed LLM Context Loss with pytest + Docker Automation

Woken up by an alarm call at 2 a.m., users were frantically complaining, "The AI suddenly doesn't remember anything I said — its memory is like a goldfish." I checked the logs and found that half of the conversation history had silently disappeared. The session was still there, no backend errors, but the context just vanished into thin air. That night I manually ran hundreds of test conversations over and over. The next day my eyes burned as if someone had thrown chili powder in them. So I finally wrote an automated test suite using pytest + Docker, caging the "memory loss" beast so it never escaped again. This post covers how we used testing to wipe out LLM memory storage issues once and for all.

Problem Breakdown: Why Does Context "Shrink"?

We integrated an LLM into our customer support chatbot and built our own memory storage layer: we embed multi-turn dialogues into vectors and store them in a vector database; subsequent requests retrieve relevant history to assemble the context. This pipeline looks simple but hides several pitfalls:

Asynchronous write race: at the end of a conversation, memory writes are async. If the next user request arrives too quickly, it may read a stale snapshot and lose some context.
Summarization/truncation strategy: to prevent token explosion, we compress or truncate history before storing. However, mishandling boundary conditions (e.g., the first message gets truncated) can split the context in two.
TTL and eviction logic: we set expiration times for memories. In the test environment, time jumps differ from production, so the cleanup job evicted memories prematurely, causing the AI to instantly forget.
ID mismatch: in multi-tenant scenarios, a session ID generation rule changed without migrating old memories; new requests read nothing at all — total chaos.

The usual approach is "manual regression": PMs click around, devs click around, fix what's broken. But context loss is often intermittent and timing-sensitive; manually clicking a hundred times might reproduce it three or four times, and a few days later it shows up in a new version wearing a different disguise. We needed an automated test suite isomorphic to the real service.

Solution Design: Why pytest + Docker?

My goal was clear: run everything automatically in CI, using real backend storage. Here's the reasoning behind the tech choices:

Why not mock the database? Because many bugs stem from real-world details like network latency, serialization, and connection pool exhaustion — a mocked world is peaceful until production blows up.
Why Docker instead of a preconfigured dev environment? To guarantee environment consistency and eliminate "works on my machine". Also, in CI we can spin up a fresh, clean container each time, leaving no leftover data from previous tests.
Why pytest? Its rich ecosystem makes managing Docker container lifecycles through fixtures a breeze. The testcontainers-python library wraps Docker into fixtures that burn after use — no more writing tons of bash scripts.

Architecture overview: Use a session-level pytest fixture to start a Redis container (similar approach for vector DBs), initialize our own MemoryStore component. Each test case simulates a complete multi-turn conversation sequence, covering concurrent writes, boundary truncation, TTL expiry, ID changes, etc. Finally, assert that the retrieved context is complete, correctly ordered, with no missing sentences.

Core Implementation: Caging Memory Storage in Tests

I'll walk through three code sections: environment setup, core test, and concurrency simulation. These are fully runnable — just tweak and use.

1. pytest fixture launches a Docker container, providing a clean Redis instance

This code solves the "brand new memory storage for each test" problem. Using testcontainers-python it automatically pulls the Redis image, starts the container, and injects connection info into the test.

# conftest.py
import pytest
from testcontainers.redis import RedisContainer
import redis

@pytest.fixture(scope="session")
def redis_container():
    """启动一个 Redis 容器，整个测试会话共用，避免频繁重启。"""
    container = RedisContainer("redis:7-alpine")
    container.start()
    yield container
    container.stop()

@pytest.fixture(scope="function")
def memory_store(redis_container):
    """每个测试函数获得一个独立的记忆存储实例，自动清空数据。"""
    client = redis.Redis(
        host=redis_container.get_container_host_ip(),
        port=redis_container.get_exposed_port(6379),
        decode_responses=True
    )
    store = MemoryStore(client, prefix="test")  # 我们的记忆存储封装
    store.flush()  # 保证测试隔离
    yield store
    store.flush()  # 清理痕迹

Key point: scope="function" ensures each test sees a clean database, thoroughly decoupling tests from each other. This is a lesson learned the hard way — we previously used module-level fixtures, and leftover keys from one test caused another to falsely pass.

2. MemoryStore core logic and the first test: context integrity

MemoryStore encapsulates dialogue history saving, retrieval, and truncation policies. The code is a bit long, so I'll show the key skeleton. The test below verifies that after appending 4 messages consecutively, the retrieved context is correctly concatenated in order, with no content missing.

# test_memory_integrity.py
from memory import MemoryStore

def test_context_retrieved_completely(memory_store: MemoryStore):
    """模拟 4 轮对话后，检索到的上下文必须包含全部消息且顺序正确。"""
    session_id = "session-abc"
    messages = [
        {"role": "user", "content": "推荐一本推理小说"},
        {"role": "assistant", "content": "《恶意》怎么样？"},
        {"role": "user", "content": "谁写的？"},
        {"role": "assistant", "content": "东野圭吾，神反转。"},
    ]

    # 模拟真实场景：每次对话结束后异步写入记忆（但我们改为同步方便测试）
    for msg in messages:
        memory_store.append(session_id, msg)

    # 检索上下文，默认返回最近 10 条
    context = memory_store.retrieve(