AI Agent 记忆存储踩坑实录：这个问题让我排查了 3 天，最终用 pytest + Docker 实现秒级回归

凌晨三点，手机炸了。同事语音里声音都在抖：线上 Agent 突然"失忆"，所有用户的对话历史全丢。我眯着眼连上 VPN 一看，好家伙------记忆存储模块的写入逻辑在处理带特殊 Unicode 字符的 metadata 时，直接抛异常把整个 batch 操作中断，连回滚都没写好，数据库连接池全僵死。

从那以后我心里就一根刺：记忆模块的测试太糙了，每次发布都像蒙眼走钢丝。于是我花了两周，用 pytest + Docker 搭了一套能在 CI 里秒级回归的自动化测试体系。踩过的坑比代码行数还多，下面全交代。

问题拆解

Agent 的记忆存储不是什么 CRUD 就完事。它至少要扛住：

• 多模态元数据：对话摘要、时间戳、token 统计、工具调用结果，一个字段的结构变更就可能让反序列化挂掉。
• 向量检索 + 过滤：按相似度搜索，再叠加时间范围、会话 ID 过滤，SQL 拼错一点就扫全表。
• 高并发追加：同一会话里消息顺序写，很容易产生竟态，导致历史条目乱序或丢失。
• 边界数据：空消息、超长摘要、list 为空的 metadata，线上脏数据比想象中丰富得多。

常规"方案"：写完代码，手动在本地起个数据库，用 curl 或 Postman 戳几下，看几条记录没问题就上线。这根本不是测试，是许愿。根因就是没有任何隔离的、可重复执行的端到端验证。单靠 mock 数据库接口，隐藏了序列化、连接池、事务隔离的真问题------那些线上爆炸的坑，mock 百分之百钓不出来。

方案设计

核心组合拳：pytest 做编排，Docker 容器提供真实但临时的数据库实例。

为什么不用 unittest？参数化、fixture 管理、第三方插件（pytest-xdist 并行、pytest-randomly 随机顺序）的生态差别，用过的都懂。最关键的是 conftest 的分层注入能力，可以把DB连接的创建、表结构初始化、隔离策略全部封装成 fixture，测试函数里干干净净。

为什么不用嵌入式的 SQLite 加向量扩展？线上跑的是 PostgreSQL + pgvector，SQLite 连索引行为都不一样，测出来的"通过"才是真正的谎话。必须用和线上一致的镜像，直接在 Docker 里起真实的 Postgres，测试完即销毁。而且还不能每个测试函数起一个全新容器------那样运行时间直接炸成半小时。用 session 级容器 + 事务回滚或临时 schema 做隔离，实测全程跑 80 多个测试用例保持在 90 秒以内。

不选 Docker Compose 前置启动，因为 testcontainers-python 库能直接用代码管理容器生命周期，测试结果和 CI 环境的宿主机无关，开发换了电脑或 CI 换了执行器，一行命令就能跑。

核心实现

这段代码解决"如何用 fixture 优雅管理 Docker Postgres 的生命周期，并确保每个测试有独立表"。

# 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),
            created_at TIMESTAMPTZ DEFAULT now()
        );
        CREATE INDEX ON memories USING ivfflat (embedding vector_cosine_ops);
    """)
    db_info.commit()
    yield db_info
    db_info.rollback()
    db_info.execute(f"DROP SCHEMA {schema} CASCADE")
    db_info.commit()
    db_info.close()

这段代码解决"如何测试边界条件：空元数据写入、超长文本插入不丢数据、向量搜索并发写"。

# test_memory_store.py
import json
import threading
import time
from memory_store import MemoryStore  # 待测模块

def test_insert_empty_metadata_should_not_crash(db):
    store = MemoryStore(db)
    memory_id = store.add_memory(session_id="s1", content="hello", metadata={})
    row = db.execute("SELECT metadata FROM memories WHERE id=%s", (memory_id,))
    assert row[0][0] == {}  # 不能存成 null 或 'null'

def test_ultra_long_content_truncation_preserves_search(db):
    store = MemoryStore(db)
    long_text = "重要内容前段。" + "废话" * 50000
    mid = store.add_memory(session_id="s2", content=long_text, metadata={"key": "val"})
    # 按 metadata 过滤能找到
    results = store.search_memories(session_id="s2", metadata_filter={"key": "val"}, limit=5)
    assert len(results) == 1
    assert results[0].id == mid

def test_concurrent_append_message_order(db):
    store = MemoryStore(db)
    errors = []
    def append_msg(seq):
        try:
            store.add_memory("s3", f"msg_{seq}", metadata={})
        except Exception as e:
            errors.append(e)
    threads = [threading.Thread(target=append_msg, args=(i,)) for i in range(20)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()
    assert not errors
    rows = db.execute("SELECT content FROM memories WHERE session_id='s3' ORDER BY id")
    contents = [r[0] for r in rows]
    # 顺序写入不应出现乱序
    assert contents == [f"msg_{i}" for i in range(20)]

这段代码解决"向量搜索质量验证：确保近似的 K 个结果真的排在前面"。

def test_vector_search_recall_at_k(db):
    store = MemoryStore(db)
    # 插入若干已知向量，target 明确靠近前3个
    base_embeddings = [[0.1]*1536, [0.2]*1536, [0.3]*1536, [0.9]*1536, [0.95]*1536]
    for i, emb in enumerate(base_embeddings):
        store.add_memory("vs", f"doc_{i}", metadata={}, embedding=emb)
    target = [0.15]*1536
    results = store.search_by_vector(target, top_k=3)
    result_texts = [r.content for r in results]
    assert "doc_0" in result_texts
    assert "doc_1" in result_texts
    assert "doc_2" in result_texts

踩坑记录

坑一：用 localhost 连不上 testcontainers 出来的数据库

现象：在 macOS 上 localhost:5432 死活连不上，报 connection refused。原因是 testcontainers 的端口映射实际绑定在所有网卡上，但 Docker Desktop 的网络转发在 macOS 上有问题，容器内的 Postgres 监听的地址绑定在容器内部 IP，而不是 0.0.0.0。后来翻 testcontainers-python 源码悟了，用 postgres.get_container_host_ip() 和 postgres.get_exposed_port(5432) 去构造连接信息，完美解决。官方文档的 quickstart 例子用的 localhost 其实只在 Linux CI 上管用。

坑二：pgvector 索引在测试里从来不生效，线上却全表扫

原因是我们建好 ivfflat 索引之后，没有用 SET ivfflat.probes = 10; ，导致执行计划选了 seqscan。但测试数据量才几十条，Postgres 查询优化器认为用不着索引，从来不走。结果测试通过，一上线万级数据直接把 CPU 打满。后来强制在测试 fixture 里执行 ANALYZE memories; 并在搜索前设置 probes，再加一个大数据量性能测试：插入 10,000 条向量，要求查询在 50ms 内返回。这才是真正端到端验证。

效果验证

测试维度	优化前（手工）	优化后（pytest + Docker）
边界条件覆盖	依赖记忆力，经常漏测	42 个专项用例，覆盖空值、超长、并发
执行时间	平均 2 小时搭建+验证	81 秒（单机）
CI 集成	不可能，每次要人肉环境	一个 pytest 命令，GitHub Actions 自动跑
线上逃生次数	3 次记忆丢失事故	0------至少两次变更被这套测试抓住回滚

可直接用的代码 / 工具

我把上面的核心 fixture 抽成了一个可复用的 pytest-memory-container 包，放在 GitHub 上：

pip install git+https://github.com/baofugege/pytest-memory-container.git
# 然后在你的 conftest.py 里:
from pytest_memory_container import postgres, db

不用再手写那堆容器管理逻辑。

#Python #后端 #AI Agent #pytest #自动化测试 #Pgvector

关于作者

一个常年和后端数据、Agent 框架死磕的实战派开发者，坚信没有自动化测试的代码就是定时炸弹。

GitHub: github.com/baofugege （上面有这篇文章提到的测试工具包）

Sponsor: github.com/sponsors/ba... ------ 如果这篇文章帮你省了三小时的排错，请我喝杯咖啡

提供服务：Python 后端性能优化 / 工具定制 / 技术咨询，联系 Telegram @baofugege