凌晨两点,我被报警电话震醒——线上Agent连续三次对话都“失忆”了,之前说好的上下文全丢,用户疯狂投诉。我眯着眼睛打开日志,发现记忆管理模块的rollback方法被一个看似无害的代码重构给搞坏了:回滚不仅撤销了错误操作,还把整个会话历史一并干掉了。最要命的是,这个问题在我们现有的单元测试里完全没暴露——因为测试每次都从空数据库开始,根本覆盖不了“脏数据回滚到上一个快照”这种跨会话场景。我花了3小时排查、手动模拟状态,才定位到根因。那一刻我才意识到:我们缺的不是测试,是能把“记忆状态”整个拍下来的快照测试。
问题拆解
我们的LLM记忆系统基于SQLite做本地持久化,每个会话维护一张表,保存对话轮次、向量摘要和工具调用记录。关键操作有两个:
save_snapshot(session_id):把当前会话的完整状态序列化到snapshots表,生成一个可回滚的检查点。rollback_to_snapshot(session_id, snapshot_id):发生异常时,根据快照重建会话表,丢弃之后的所有变更。
这套机制平时跑得好好的,但那次重构我改动了回滚逻辑中的事务边界,导致rollback执行后,conversation表被重建,但snapshots表自身也被意外清空——下一次回滚就再也找不到更早的检查点了。
常规单元测试为什么抓不住这个bug?因为测试套路是:
def test_rollback():
db = create_in_memory_db()
db.save_snapshot("s1")
db.rollback_to_snapshot("s1", ...)
assert db.get_conversation("s1") == expected
这一切都发生在一个进程、一个临时数据库里。但线上问题是:进程A保存快照后退出,进程B重新打开同一数据库文件并执行回滚。文件级的持久化状态、WAL日志的合并、甚至不同连接对snapshots表的可见性,这些统统没测到。说白了,我们测了“逻辑”,却没测“存储”。
方案设计
我决定引入**快照测试(snapshot testing)**的思想,但不用常见的文本快照,而是把SQLite数据库文件本身当作不可变快照。
选型对比:
- pytest-snapshot:只做文本/JSON快照,不适合二进制或复杂状态比对。
- 直接用pytest的
tmp_path+手动比对:灵活,但每次自己写比对逻辑容易遗漏字段。 - 基于文件hash + SQLite内对比:先对数据库文件计算sha256作为整体快照,同时可选导出关键表的内容做可读差异对比。这样做既能快速发现“哪里变了”,又能保留细粒度调试能力。
架构思路:pytest的conftest.py提供一个snapshot_db fixture,它负责:
- 在测试开始前,检查是否存在基线快照文件(比如
tests/snapshots/memory_test.sqlite)。 - 如果没有,则自动生成(
--snapshot-update模式),测试直接pass。 - 如果有,则执行测试操作后,对被测数据库文件做sha256,并与基线文件的sha256对比。如果不一致,测试失败并输出差异提示。
这样,我们的测试就真正模拟了“跨进程、跨连接”的持久化效果——每个测试用例拿到的都是一个独立拷贝的数据库文件,执行完操作后比较整个文件状态是否与预期完全一致。
核心实现
1. 先搭一个可持久化的记忆管理器
这段代码解决“我们到底要测什么”的问题。MemoryManager封装了SQLite的连接、快照保存和回滚,是线上真实使用的简化版。
# memory_manager.py
import sqlite3
import uuid
from datetime import datetime, timezone
class MemoryManager:
def __init__(self, db_path: str):
self.db_path = db_path
self._init_tables()
def _get_conn(self) -> sqlite3.Connection:
conn = sqlite3.connect(self.db_path)
conn.execute("PRAGMA journal_mode=WAL")
conn.row_factory = sqlite3.Row
return conn
def _init_tables(self):
with self._get_conn() as conn:
conn.executescript("""
CREATE TABLE IF NOT EXISTS conversations (
session_id TEXT NOT NULL,
turn INTEGER NOT NULL,
role TEXT NOT NULL,
content TEXT NOT NULL,
PRIMARY KEY (session_id, turn)
);
CREATE TABLE IF NOT EXISTS snapshots (
snapshot_id TEXT PRIMARY KEY,
session_id TEXT NOT NULL,
created_at TEXT NOT NULL,
state_json TEXT NOT NULL
);
""")
def add_message(self, session_id: str, role: str, content: str):
with self._get_conn() as conn:
turn = conn.execute(
"SELECT COALESCE(MAX(turn), 0) + 1 FROM conversations WHERE session_id = ?",
(session_id,)
).fetchone()[0]
conn.execute(
"INSERT INTO conversations VALUES (?, ?, ?, ?)",
(session_id, turn, role, content)
)
def save_snapshot(self, session_id: str) -> str:
snapshot_id = str(uuid.uuid4())
with self._get_conn() as conn:
rows = conn.execute(
"SELECT turn, role, content FROM conversations WHERE session_id = ? ORDER BY turn",
(session_id,)
).fetchall()
import json
state = [dict(r) for r in rows]
conn.execute(
"INSERT INTO snapshots VALUES (?, ?, ?, ?)",
(snapshot_id, session_id, datetime.now(timezone.utc).isoformat(), json.dumps(state))
)
return snapshot_id
def rollback_to_snapshot(self, session_id: str, snapshot_id: str):
with self._get_conn() as conn:
row = conn.execute(
"SELECT state_json FROM snapshots WHERE snapshot_id = ?",
(snapshot_id,)
).fetchone()
if not row:
raise ValueError("Snapshot not found")
import json
state = json.loads(row["state_json"])
conn.execute("DELETE FROM conversations WHERE session_id = ?", (session_id,))
for msg in state:
conn.execute(
"INSERT INTO conversations VALUES (?, ?, ?, ?)",
(session_id, msg["turn"], msg["role"], msg["content"])
)
2. 快照fixture:将数据库变成可验证的“照片”
这段代码解决“怎么自动比对整个数据库文件”的问题。fixture会复制一个干净的模板数据库,测试操作完成后,对最终数据库文件取sha256,与基线快照文件的sha256比较。
# conftest.py
import pytest
import shutil
import hashlib
import os
from pathlib import Path
SNAPSHOT_DIR = Path(__file__).parent / "snapshots"
def _db_checksum(path: str) -> str:
with open(path, "rb") as f:
return hashlib.sha256(f.read()).hexdigest()
@pytest.fixture
def snapshot_db(request, tmp_path):
"""
返回一个 MemoryManager 实例,其背后是一个可落盘的数据库文件。
测试结束后,自动比对数据库文件与基线快照。
"""
# 基线快照路径:测试函数名+.sqlite
snapshot_name = request.node.name + ".sqlite"
snapshot_path = SNAPSHOT_DIR / snapshot_name
# 在临时目录创建被测数据库(先复制基线,如果存在)
db_tmp = tmp_path / "test.db"
if snapshot_path.exists():
shutil.copy(snapshot_path, db_tmp)
else:
# 首次运行:创建空数据库,后续由 --snapshot-update 模式生成基线
pass
manager = MemoryManager(str(db_tmp))
yield manager
# 检查是否需要更新快照
update_snapshots = request.config.getoption("--snapshot-update", default=False)
manager._get_conn().close() # 确保 WAL 刷盘
if update_snapshots:
SNAPSHOT_DIR.mkdir(parents=True, exist_ok=True)
shutil.copy(db_tmp, snapshot_path)
else:
if not snapshot_path.exists():
pytest.fail(f"Snapshot file {snapshot_path} missing. Run with --snapshot-update to create.")
expected_checksum = _db_checksum(str(snapshot_path))
actual_checksum = _db_checksum(str(db_tmp))
if actual_checksum != expected_checksum:
pytest.fail(
f"Database snapshot mismatch.\n"
f"Expected: {expected_checksum}\n"
f"Actual: {actual_checksum}\n"
f"Tip: run with --snapshot-update to regenerate."
)
3. 写测试用例:验证保存后回滚不会丢数据
这段代码模拟了“存快照→加新消息→回滚”的关键路径,并依靠fixture自动核对数据库完整性。
# test_memory.py
import pytest
def test_rollback_preserves_earlier_snapshots(snapshot_db):
session = "sess_01"
# 初始对话
snapshot_db.add_message(session, "user", "hello")
snapshot_db.add_message(session, "assistant", "hi there")
snap_id = snapshot_db.save_snapshot(session)
# 再次对话并产生新快照
snapshot_db.add_message(session, "user", "how are you?")
snapshot_db.add_message(session, "assistant", "I'm fine")
snap_id2 = snapshot_db.save_snapshot(session)
# 回滚到第一个快照
snapshot_db.rollback_to_snapshot(session, snap_id)
# 此时数据库状态应等于第一次快照后的状态
# 快照fixture会在测试结束后做全库比对,如果意外删除了 snapshots 表内容,文件checksum必变
# 我们还可以做轻量断言辅助定位
rows = snapshot_db._get_conn().execute(
"SELECT turn, content FROM conversations WHERE session_id=? ORDER BY turn", (session,)
).fetchall()
assert len(rows) == 2
assert rows[1]["content"] == "hi there"
此刻,任何破坏回滚逻辑的修改——比如错误地清空了snapshots表——都会被文件级的checksum比对抓出来,因为清空操作改变了数据库页,即使最终会话数据看起来“对”了,整个文件也不再匹配基线。
踩坑记录
坑1:WAL模式下的“影子数据库”
现象:第一次跑快照测试,checksum总是不一样,即使逻辑完全正确。打开基线文件对比,大小有时差几十KB。原因:SQLite在WAL模式下,主数据库文件不会立即写入最新数据,而是先写-wal文件,只有checkpoint之后才会合并。我们的fixture在测试结束时直接对test.db读hash,可能读到一个尚未合并WAL的状态。
解决:在finish阶段,显式关闭连接(manager._get_conn().close())并在关闭前执行PRAGMA wal_checkpoint(TRUNCATE)。但最简单的方式是关闭所有连接后,SQLite会自动做checkpoint。我们的fixture在关闭连接后才计算hash,保证了文件完整性。这是官方文档没明说、实际调试时才会发现的细节。
坑2:快照文件没进版本控制,CI全红
现象:本地跑全绿,推到CI上直接挂,报Snapshot file missing。因为我们团队有人忘记git add新生成的快照文件,而CI环境从不运行--snapshot-update。 解决:在CI脚本里加了一个防御性检查——如果tests/snapshots/下有未跟踪文件,直接失败并提示开发者提交。同时在README里写清楚:任何修改记忆管理模块的PR,必须附带更新后的快照文件。这也倒逼团队养成了“改持久化先看快照”的习惯。
效果验证
以前手动回归整个记忆回滚流程,需要分3个场景、来回构造SQLite文件,最快也要30分钟,还容易漏测。现在:
- 测试覆盖场景数:1个 → 7个(正常回滚、回滚到空快照、回滚不存在的id、并发保存、WAL积压后回滚等)
- 单次回归耗时:30分钟 → 12秒
- Bug发现率:线上记忆丢失事件后再无同类故障,两次重构均被快照测试挡在PR阶段
快照比对让“持久化是否变脏”变成二进制判断,再也不用靠人肉看日志猜状态了。
直接可用的代码
# 一行命令更新全部快照基线,开发者本地执行
pytest --snapshot-update
把这套conftest.py和SNAPSHOT_DIR放入你的项目,接入任何依赖文件状态的模块,立刻获得跨进程的持久化验证能力。
#Python #LLM #测试 #后端 #SQLite
关于作者
一个专啃后端硬骨头的实战派开发者,擅长把内存、存储、并发这些“玄学”问题变成代码就能验证的确定性方案。
GitHub: github.com/baofugege — 本项目的完整快照测试框架模板即将开源。
Sponsor: github.com/sponsors/ba… — 如果这篇文章帮你省下3小时排查时间,请我喝杯咖啡。
提供服务:Python后端性能优化 / 工具定制 / 技术咨询,联系 Telegram @baofugege