凌晨两点被报警电话叫醒——线上 AI 助手突然“失忆”了，用户问“上次我们聊到哪儿了”，它回了一句“请问有什么可以帮您？”。查日志发现是向量数据库的升级脚本改了写入路径，老记忆全部写入新 collection，但检索时还去读旧的。手工回归一遍所有记忆场景至少两个小时，还不保证全覆盖。那次之后我决心把这坨手工测试彻底干掉，用 pytest + Docker 造一套自动化验证流水线。现在所有记忆存储变更，3 分钟 15 条用例全过，漏测直接清零。
问题拆解：为什么 AI 记忆存储的一致性这么难测
AI 应用的“记忆”不是简单的 SQL 行，它涉及文本摘要 → embedding 向量 → 向量数据库写入 → 相似度检索 → 带上下文拼接的完整链路。任何一个节点出问题，都可能让助手“失忆”或“记串线”。我们团队用的是 Chroma 作为高性能向量存储，配合自定义的 MemoryManager 做记忆的增、删、模糊检索。日常迭代中经常要改 embedding 模型、调整分块策略、甚至升级 Chroma 版本。
原本的测试方式是：改完代码，在本地手动拉起一个 Chroma 实例，用 curl 或几段临时代码塞几条记忆，再肉眼比对检索结果是否一致。这种方式有三个致命问题：
- 状态污染严重——上一个用例的残留数据会影响下一个，经常要手动清 collection，一不留神就“怎么刚才还过的测试突然崩了？”
- 覆盖粒度靠人脑——15 个场景，跑着跑着就忘了第 9 个到底测没测过，全靠纸笔记 checklist。
- 回归成本巨大——每次发版前都要重跑一遍，2 小时起跳，CI 根本没法集成。
更糟的是，单元测试里用 mock 把 Chroma 的 Client 换掉，完全避开了真实的网络 I/O、embedding 计算和向量比对，等于自欺欺人。我们要的是在真实环境里跑断言，而不是对着假数据测逻辑。
方案设计：为什么选 pytest + Docker，而不是其他组合
我需要的方案满足三点：
- 环境可销毁：每次测试都是全新的 Chroma，不留历史数据。
- 全链路真实：真正调用 embedding 模型，真正写入磁盘/内存索引，真正计算余弦距离。
- 能进 CI：命令行一把梭，开发者本地能跑，CI runner 也能跑，且速度不超过 5 分钟。
为什么不选 mock 单元测试？上面说了，mock 掉的 I/O 不会告诉你 embedding 模型跟 Chroma 的 dim 不匹配，也不会暴露索引重建后的检索差异。
为什么不用 全栈 E2E ？启动整个 AI 应用 + LLM 服务太重，动不动 10 分钟以上，不适合高频回归。
最终我选了 pytest + testcontainers + chromadb 这一套：
testcontainers用代码管理 Docker 容器，不需要额外写 docker-compose，容器生命周期跟 fixture 绑定，pytest 退出自动销毁。chromadb的Client直接连接容器内部 HTTP 端口，真实的客户端体验。- 每个测试用例前通过 fixture 自动创建独立 collection，用例结束直接删除，杜绝污染。
架构极度简单：pytest fixture 启动 Chroma Docker → 返回 Client → 测试函数执行记记存储/检索 → 断言一致性 → 自动清理。没有第三方 mock，没有中间件。
核心实现：把测试写成活文档
1. 用 fixture 管理 Chroma 容器生命周期
这段代码解决“怎么让数据库自己活过来，测完自己死掉”的问题。用 testcontainers 的 GenericContainer 拉取 Chroma 镜像，并等待服务就绪。
# conftest.py
import pytest
from testcontainers.core.container import GenericContainer
from testcontainers.core.waiting_utils import wait_for_logs
import chromadb
from chromadb.config import Settings
@pytest.fixture(scope="session")
def chroma_container():
"""启动 Chroma 容器，返回容器对象，session 级复用"""
container = (
GenericContainer("chromadb/chroma:0.4.22")
.with_exposed_ports(8000)
)
container.start()
# 等待日志确认服务就绪，避免客户端握手失败
wait_for_logs(container, "Uvicorn running on http://0.0.0.0:8000")
yield container
container.stop()
2. 夹具提供隔离的 Client 与 Collection
这个 fixture 在每个测试函数前自动销毁旧 collection、创建新 collection，保证用例之间零干扰。
@pytest.fixture
def chroma_client(chroma_container):
"""返回连接容器内 Chroma 的 Client"""
host = chroma_container.get_container_host_ip()
port = chroma_container.get_exposed_port(8000)
return chromadb.Client(Settings(
chroma_api_impl="rest",
chroma_server_host=host,
chroma_server_http_port=port
))
@pytest.fixture
def memory_collection(chroma_client, request):
"""
为每个测试函数创建独立 collection，测试结束直接删除。
collection 名使用测试函数名，方便问题回溯。
"""
collection_name = f"test_{request.node.name}"
collection = chroma_client.create_collection(collection_name)
yield collection
chroma_client.delete_collection(collection_name)
3. 测试记忆存储与检索一致性——带元数据过滤
下面的测试验证“写入一对对话记忆，并用 user_id 过滤检索，返回的文档和元数据必须完全匹配”。这里会真实触发 Chroma 的 embedding API 调用（如果使用默认的 all-MiniLM-L6-v2 嵌入模型）。
# test_memory_consistency.py
import uuid
def test_memory_insert_and_retrieve_by_metadata(memory_collection):
"""测试记忆写入后，按元数据过滤完整取出"""
user_id = str(uuid.uuid4())
doc_text = "用户喜欢喝拿铁，不加糖"
meta = {"user_id": user_id, "category": "preference"}
# 写入记忆，Chroma 自动生成 embedding
memory_collection.add(
documents=[doc_text],
metadatas=[meta],
ids=[f"mem_{user_id}_1"]
)
# 用 user_id 过滤检索
results = memory_collection.query(
query_texts=["咖啡偏好"],
where={"user_id": user_id},
n_results=1,
include=["documents", "metadatas"]
)
assert len(results["ids"][0]) == 1
# 验证文档内容不会被“压缩”或截断
assert results["documents"][0][0] == doc_text
# 验证元数据原路返回，尤其 user_id 不能丢失
assert results["metadatas"][0][0]["user_id"] == user_id
4. 记忆更新后旧向量必须“失活”
很多记忆模块的 bug 就出在“更新了文档但旧向量还在索引里”，导致问同一个问题同时召回新老两条记录。下面测试验证：更新记忆后，用相同的语义查询只能返回新内容。
def test_memory_update_removes_old_vector(memory_collection):
"""验证更新记忆后，旧向量不会被检索到"""
memory_collection.add(
documents=["用户住在深圳"],
metadatas=[{"user_id": "u1"}],
ids=["mem_u1_addr"]
)
# 更新成新城市
memory_collection.update(
ids=["mem_u1_addr"],
documents=["用户搬到了北京"],
metadatas=[{"user_id": "u1"}]
)
res = memory_collection.query(
query_texts=["住在哪里"],
n_results=1
)
# 只应返回新文档，绝不能出现深圳
assert "北京" in res["documents"][0][0]
assert "深圳" not in res["documents"][0][0]
踩坑记录：官方文档没说清的事
坑 1：容器启动的“虚假就绪”
现象：chroma_client 创建成功，但第一个 create_collection 直接抛 ConnectionError。原因：Chroma 容器在 HTTP 端口暴露后，内部 embedding 模型后台加载还没完成，此时任何 API 调用都会失败。解决：用 wait_for_logs 捕获特定日志 "Uvicorn running on http://0.0.0.0:8000" 后再返回容器对象，确保真正就绪。网上很多教程直接 time.sleep(5)，在 CI 资源紧张时完全不可靠。
坑 2：Chroma 默认按余弦距离检索，阈值过滤的误解
我们想验证“相关度低的记忆不会被返回”，于是写了 assert results["distances"][0][0] < 0.3。本地跑全绿，CI 上 Arch Linux 的 Docker 宿主偶尔失败。排查发现 Chroma 的 query 默认不带阈值，n_results 是硬返回量，和距离无关。就算所有向量相似度 0.8，它也会硬塞给你 N 条。要模拟真实业务过滤，必须自己在应用层设置距离阈值，而不是指望 Chroma 替你裁剪。这个行为在文档里藏在 query 的高级参数部分，开始没注意。
效果验证：从人工 2 小时到自动化 3 分钟
| 对比维度 | 手工测试 | pytest + Docker 自动化 |
|---|---|---|
| 全场景覆盖 | 15 条用例容易遗漏，经常漏测 3-4 条 | 100% 覆盖，pytest 输出清晰列出 |
| 单次耗时 | ~120 分钟（含环境准备、手动清理） | 3 分 12 秒（包含容器启动 & 15 个测试） |
| 可重复性 | 环境污染导致偶发不可复现 | 每次全新 collection，结果确定 |
| CI 集成 | 无，手工回归 | 一键 make test-memory，push 自动跑 |
| 漏测风险 | 高，上次因漏测记忆分块逻辑，线上白屏 | 0，全部自动化断言 |
在 GitHub Actions 上跑也不用额外配置 Docker 环境，因为 testcontainers 利用 host Docker 套接字，Runner 自带。内存占用峰值仅 500MB，免费 tier 就能跑。
直接拿去用的命令 & 模板
如果你也用 Chroma 做记忆存储，可以直接复制以下 fixture 到 conftest.py，然后运行：
pip install pytest testcontainers chromadb
pytest tests/ -m memory_consistency -v --tb=short
想跑在 CI 里的 GitHub Action 配置（精简版）：
- name: Run memory consistency tests
run: |
pip install .[dev]
pytest tests/memory -v --tb=short
不需要额外启动数据库服务，代码自己管。
#Python #AI工程化 #测试自动化 #Docker #Chroma
关于作者
一个长期跟向量数据库和 LLM 应用打交道的后端/架构开发者，坚信“不能自动化的测试都会回来咬你一口”。
GitHub: github.com/baofugege
Sponsor: github.com/sponsors/ba… — 如果这篇文章省下了你加班时间，请我喝杯咖啡。
提供服务：Python 后端性能优化 / 自动化测试工具定制 / 技术咨询，联系 Telegram @baofugege
