At 3 a.m., my phone exploded with notifications. A colleague's voice trembled over a voice call: the production AI agent suddenly "lost its memory" — all user conversation histories vanished. I squinted, VPNed in, and wow — the memory storage module's write logic threw an exception when handling metadata with special Unicode characters, aborting the entire batch operation. Rollback wasn't even properly implemented, and the database connection pool froze solid.
Since then, it's been a thorn in my side: the memory module's testing was too sloppy. Every deployment felt like walking a tightrope blindfolded. So I spent two weeks building an automated test system with pytest and Docker that can run regression tests in seconds in CI. I stumbled into more pitfalls than lines of code — here's the full story.
Breaking Down the Problem
Agent memory storage is far more than simple CRUD. It must at least withstand:
- Multimodal metadata: Conversation summaries, timestamps, token counts, tool call results — a single field's structural change can break deserialization.
- Vector search + filtering: Similarity search combined with time range and session ID filters. A single SQL typo and you're scanning the entire table.
- High-concurrency appends: Sequential writes within the same session easily cause race conditions, leading to out-of-order or missing history entries.
- Edge-case data: Empty messages, super-long summaries, metadata with empty lists — production dirty data is far richer than you'd imagine.
The usual "approach": write the code, manually spin up a local database, poke it with curl or Postman, check a few records look fine, then ship it. That's not testing — it's wishful thinking. The root cause is the absence of any isolated, reproducible end-to-end verification. Relying solely on mocking database interfaces hides the real issues of serialization, connection pooling, and transaction isolation — the kind of problems that blow up in production. Mocks will never catch them, 100%.
Solution Design
The core combination: pytest for orchestration, Docker containers to provide real but temporary database instances.
Why not unittest? The ecosystem differences — parameterization, fixture management, third-party plugins (pytest-xdist for parallelism, pytest-randomly for random order) — are obvious to anyone who's used both. Most importantly, conftest’s layered injection ability lets you encapsulate DB connection creation, schema initialization, and isolation strategies entirely within fixtures, leaving test functions clean.
Why not an embedded SQLite with vector extensions? Production runs on PostgreSQL + pgvector. SQLite behaves differently even with indexing — a "pass" there would be a true lie. You must use the same image as production, spinning up a real Postgres in Docker and tearing it down after tests. And you can't start a new container per test function — that would blow up runtime to half an hour. Using session-scoped containers plus transaction rollback or temporary schemas for isolation, I kept 80+ test cases running in under 90 seconds.
No Docker Compose pre-startup: the testcontainers-python library manages container lifecycles directly in code. Test results become independent of the CI host environment. Switch laptops or CI runners, and it's still a single command to execute.
Core Implementation
This snippet shows how to elegantly manage a Docker Postgres lifecycle with fixtures and ensure each test gets its own isolated tables.
# conftest.py
import pytest
from testcontainers.postgres import PostgresContainer
import psycopg2
from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT
import uuid
class MemoryDB:
"""封装连接,避免直接暴露 psycopg2 游标"""
def __init__(self, host, port, db, user, password):
self.conn = psycopg2.connect(
host=host, port=port, dbname=db, user=user, password=password
)
self.conn.set_isolation_level(ISOLATION_LEVEL_READ_COMMITTED)
def execute(self, sql, params=None):
with self.conn.cursor() as cur:
cur.execute(sql, params)
if cur.description:
return cur.fetchall()
def commit(self):
self.conn.commit()
def rollback(self):
self.conn.rollback()
def close(self):
self.conn.close()
# session 级容器,只起一次
@pytest.fixture(scope="session")
def postgres():
postgres = PostgresContainer("pgvector/pgvector:pg16")
postgres.start()
yield postgres
postgres.stop()
# function 级隔离:每个测试建一个随机 schema,测试结束直接 drop cascade
@pytest.fixture
def db(postgres):
schema = f"test_{uuid.uuid4().hex[:8]}"
db_info = MemoryDB(
host=postgres.get_container_host_ip(),
port=postgres.get_exposed_port(5432),
db=postgres.dbname,
user=postgres.username,
password=postgres.password,
)
db_info.execute(f"CREATE SCHEMA {schema}")
db_info.execute(f"SET search_path TO {schema}")
# 建表及向量扩展
db_info.execute("""
CREATE TABLE memories (
id SERIAL PRIMARY KEY,
session_id TEXT NOT NULL,
content TEXT,
metadata JSONB DEFAULT '{}',
embedding vector(1536),
Top comments (0)