Cutting AI Agent Memory Testing from 40 Minutes to 3 with pytest + Docker

Last week I added a “remember user dietary preferences” feature to an AI agent. It worked fine in staging, but after deployment a user noticed it had “amnesia” by the third conversation. Checking the logs, writes returned success but recall came back empty. It turned out the Redis key expiration was being overwritten by an inconspicuous renewal logic. Manual verification took over 40 minutes per run, so I decided to automate this once and for all with pytest + Docker.

Problem breakdown

An AI agent’s memory storage is usually not just simple database reads and writes. It involves at least three actions: writing memories, recalling them by criteria, and expiring them by TTL. The problem is these three actions don’t follow the same path: a write may serialize successfully, but recall may still get nothing because of key structure or expiration policy.

We found two root causes: first, we used HSET to pack all memories of the same agent into one hash key, and after each write we ran EXPIRE on the whole key. As a result, frequent writes kept renewing the key lifetime, so some memories never expired. Second, the Redis version in the test environment was different from production, so local manual testing couldn’t reproduce the expiration behavior at all. The usual approaches—mocking Redis or writing unit tests—can only verify code logic. They can’t expose real Redis TTL semantics or Docker environment differences. That’s why the tests needed to run against a real Redis container and clean up automatically after each run.

Solution design

The core stack: pytest handles test orchestration, Docker via testcontainers starts a real Redis 7.2 instance, and redis-py is the client.

The reason for not mocking Redis is simple: mocks hide ordering issues between SET and EXPIRE, and they also swallow environment issues like “the service returns a connection before it’s actually listening on the port.” I didn’t manually manage docker-compose because adding manual up/down steps before every test run is tedious, and CI easily ends up with leftover dirty containers. testcontainers creates and destroys containers with the test lifecycle, which fits this kind of one-off verification perfectly.

Architecturally, we wrapped memory storage in a MemoryStore class and split each memory into its own Redis key so TTL can be controlled independently. In the tests I used a 1-second TTL to simulate expiration without long waits. Assertions cover four scenarios: write-recall consistency, expiration consistency, TTL refresh on overwrite, and batch writes.

Core implementation

This snippet handles the write, recall, and expiration encapsulation for memory storage. Each memory gets its own key to prevent TTLs from stepping on each other.

import json
from typing import Optional
import redis

class MemoryStore:
    """AI Agent 记忆存储：基于 Redis 的写入、召回、过期实现"""

    def __init__(self, host: str, port: int, db: int = 0):
        self.client = redis.Redis(host=host, port=port, db=db, decode_responses=True)

    def remember(self, agent_id: str, memory_id: str, data: dict,
                 ttl_seconds: int = 3600) -> None:
        key = f"agent:{agent_id}:memory:{memory_id}"
        # SET 本身支持 ex 参数，一条命令完成写入 + TTL，避免先写后 expire 的窗口问题
        self.client.set(key, json.dumps(data, ensure_ascii=False), ex=ttl_seconds)

    def recall(self, agent_id: str, memory_id: str) -> Optional[dict]:
        key = f"agent:{agent_id}:memory:{memory_id}"
        raw = self.client.get(key)
        if raw is None:
            return None
        return json.loads(raw)

This fixture ensures the test environment uses a real Redis and waits until the service is actually available.

import time
import pytest
import redis
from testcontainers.redis import RedisContainer

@pytest.fixture(scope="session")
def redis_container():
    # 启动真实 Redis 7.2 容器，避免 mock 带来的假阳性
    with RedisContainer("redis:7.2-alpine") as container:
        host = container.get_container_host_ip()
        port = container.get_exposed_port(6379)

        # 显式等待服务可用，避免容器端口刚映射出来但 Redis 还没监听
        client = redis.Redis(host=host, port=port, decode_responses=True)
        for _ in range(20):
            try:
                client.ping()
                break
            except redis.ConnectionError:
                time.sleep(0.2)
        else:
            raise RuntimeError("Redis container not ready")

        yield host, port

These test cases verify write-then-recall, expiry making recall impossible, TTL refresh on overwrite, and batch writes.

import time
import pytest
from memory_store import MemoryStore


@pytest.fixture
def store(redis_container):
    host, port = redis_container
    store = MemoryStore(host=host, port=port)
    yield store
    store.client.flushdb()  # 每个测试后清空，保证隔离


def test_write_then_recall_returns_same_data(store):
    memory = {"preference": "low_salt", "allergy":