把AI聊天机器人记忆存储测试从3小时压到5分钟，我用Pytest + Docker搭了一套自动化回归

凌晨一点，被 Slack 告警震醒：有用户投诉"它突然失忆了"------明明上一轮告诉 AI 我是杭州人，下一轮再问"我家乡是哪儿"，机器人一本正经说"我不知道"。一看最近发布，果然是记忆存储模块改动了序列化逻辑，历史会话加载路径被悄悄掰弯。手动回归一轮要 3 小时：分别搭 Redis、Postgres 环境，模拟多轮对话，切用户，看记忆恢复......那次之后我决定：这种体力活，必须一次性自动化干掉。

问题到底痛在哪

AI 聊天机器人的记忆存储，不是存一条 KV 那么简单。它涉及：

• 多存储后端：Redis 存短期，Postgres 存长期，甚至混合架构
• 记忆生命周期：读写、过期、合并、回滚
• 上下文与用户隔离：同一用户的多 session 串记忆，不同用户绝不串号
• 发布变更敏感：序列化方式 / Schema / TTL 逻辑一改，就可能静默丢记忆

常规方案是单元测试 mock 掉存储层，但 mock 根本测不出真实序列化与反序列化的坑------很多线上故障都是 mock 绿油油，上线火葬场。而手动回归一遍需要：拉环境 → 构造对话数据 → 发版本 → 查记忆 → 再构造 → 再查。半小时打底，全场景 3 小时还容易漏。

方案设计：把真实中间件变成测试的"即抛型"环境

核心思路：用 Docker 容器提供真实 Redis/Postgres，用 Pytest 做参数化驱动，把记忆逻辑变成可重复的回归用例。

为什么不选别的：

• 不用 fakeredis / testcontainers-python 的内存假 Redis？因为序列化/反序列化、集群模式下 Lua 脚本行为，假库还原不了，徒增虚假信心。
• 不用集成到 CI 的远程共享环境？多个测试并行会相互污染，且成本高、速度慢。
• 就用 Docker + pytest-docker 插件？插件的 session scope 管理有点坑（后面会讲），最终我选了 pytest 的 fixture + docker-compose 手动生命周期控制，灵活且稳定。

架构模型：docker-compose.yml 定义 redis 与 postgres 服务；conftest.py 里 session 级别 fixture 负责 docker-compose up，测试结束 down；测试用例用 parametrize 覆盖不同存储 backend、不同对话轮次、不同数据大小，一键回归记忆读写、隔离、迁移到底正不正常。

核心实现：一步步搭出可跑的自动化回归

1. 用 docker-compose 定义随时可销毁的真实中间件

这段配置解决"不同环境数据库不一致"的问题------开发、CI、本地全部用同一套镜像版本。

# docker-compose.yml
version: "3.8"
services:
  redis:
    image: redis:7-alpine
    ports:
      - "6379"
    # 不暴露固定端口，避免冲突
  postgres:
    image: postgres:15-alpine
    environment:
      POSTGRES_USER: test
      POSTGRES_PASSWORD: test
      POSTGRES_DB: memory_test
    ports:
      - "5432"

2. conftest.py 里用 fixture 管理容器生命周期，并注入动态连接信息

这段代码解决"每次跑测试都要手动启动服务"的痛点，且端口动态分配，不怕并行冲突。

# conftest.py
import pytest
import docker
import time
import os

# 使用 module scope，一个测试模块启动一次，平衡速度和隔离性
@pytest.fixture(scope="module")
def docker_services():
    client = docker.from_env()
    # 使用 docker-compose 启动，isolated project name 防冲突
    project_name = f"memtest_{os.getpid()}"
    compose_file = os.path.join(os.path.dirname(__file__), "docker-compose.yml")
    # docker-compose CLI 方式启动（也可以用 Compose API，但 CLI 更通用）
    os.system(f"docker-compose -p {project_name} -f {compose_file} up -d")
    # 等待 Redis/Postgres ready（健康检查轮询）
    redis_container = None
    pg_container = None
    for c in client.containers.list(filters={"name": project_name}):
        if "redis" in c.name:
            redis_container = c
        elif "postgres" in c.name:
            pg_container = c
    
    # 轮询 Redis ready
    for _ in range(30):
        exit_code, _ = redis_container.exec_run("redis-cli ping")
        if b"PONG" in exit_code:
            break
        time.sleep(0.5)
    else:
        raise RuntimeError("Redis 没起来")
    
    # 等待 Postgres 接受连接
    for _ in range(30):
        exit_code, _ = pg_container.exec_run("pg_isready -U test")
        if b"accepting" in exit_code:
            break
        time.sleep(0.5)
    else:
        raise RuntimeError("Postgres 没起来")
    
    # 拿到动态端口
    redis_port = int(redis_container.attrs["NetworkSettings"]["Ports"]["6379/tcp"][0]["HostPort"])
    pg_port = int(pg_container.attrs["NetworkSettings"]["Ports"]["5432/tcp"][0]["HostPort"])
    
    yield {
        "redis": {"host": "localhost", "port": redis_port, "db": 0},
        "postgres": {
            "host": "localhost",
            "port": pg_port,
            "user": "test",
            "password": "test",
            "dbname": "memory_test",
        },
    }
    
    # teardown: 清理所有容器
    os.system(f"docker-compose -p {project_name} -f {compose_file} down -v")

踩坑预警：用 docker-compose CLI 而非 Python SDK 的 Compose API，是因为后者在 CI 的 alpine 镜像里常缺 lib，直接调 CLI 更糙快稳。

3. 抽象 MemoryBackend 接口，写回归用例覆盖写、读、隔离与恢复

这段代码解决"不同存储后端切换测试"的问题------参数化注入 Redis 和 Postgres 实现，同一套用例全跑。

# test_memory_regression.py
import pytest
from abc import ABC, abstractmethod
import redis
import psycopg2
import json
from contextlib import contextmanager

# ---------- 统一接口 ----------
class MemoryBackend(ABC):
    @abstractmethod
    def save_memory(self, user_id: str, session_id: str, messages: list):
        pass

    @abstractmethod
    def load_memory(self, user_id: str, session_id: str) -> list:
        pass

    @abstractmethod
    def clear(self):
        pass

class RedisMemoryBackend(MemoryBackend):
    def __init__(self, host, port, db=0):
        self.client = redis.Redis(host=host, port=port, db=db, decode_responses=True)
    
    def save_memory(self, user_id, session_id, messages):
        key = f"mem:{user_id}:{session_id}"
        self.client.set(key, json.dumps(messages))
        # 模拟 TTL 策略
        self.client.expire(key, 300)
    
    def load_memory(self, user_id, session_id):
        key = f"mem:{user_id}:{session_id}"
        data = self.client.get(key)
        return json.loads(data) if data else []
    
    def clear(self):
        self.client.flushdb()

class PostgresMemoryBackend(MemoryBackend):
    def __init__(self, host, port, user, password, dbname):
        self.conn = psycopg2.connect(host=host, port=port, user=user, password=password, dbname=dbname)
        self._init_table()
    
    def _init_table(self):
        with self.conn.cursor() as cur:
            cur.execute("""
                CREATE TABLE IF NOT EXISTS memory (
                    user_id TEXT,
                    session_id TEXT,
                    messages JSONB,
                    updated_at TIMESTAMP DEFAULT NOW(),
                    PRIMARY KEY (user_id, session_id)
                )
            """)
            self.conn.commit()
    
    def save_memory(self, user_id, session_id, messages):
        with self.conn.cursor() as cur:
            cur.execute(
                "INSERT INTO memory (user_id, session_id, messages) VALUES (%s,%s,%s) "
                "ON CONFLICT (user_id, session_id) DO UPDATE SET messages = EXCLUDED.messages, updated_at = NOW()",
                (user_id, session_id, json.dumps(messages))
            )
            self.conn.commit()
    
    def load_memory(self, user_id, session_id):
        with self.conn.cursor() as cur:
            cur.execute("SELECT messages FROM memory WHERE user_id=%s AND session_id=%s", (user_id, session_id))
            row = cur.fetchone()
            return json.loads(row[0]) if row else []
    
    def clear(self):
        with self.conn.cursor() as cur:
            cur.execute("DELETE FROM memory")
            self.conn.commit()

# ---------- 测试用例 ----------
@pytest.fixture(params=["redis", "postgres"])
def backend(request, docker_services):
    if request.param == "redis":
        cfg = docker_services["redis"]
        b = RedisMemoryBackend(cfg["host"], cfg["port"], cfg["db"])
    else:
        cfg = docker_services["postgres"]
        b = PostgresMemoryBackend(cfg["host"], cfg["port"], cfg["user"], cfg["password"], cfg["dbname"])
    yield b
    b.clear()  # 每个用例结束都清空，保证隔离

class TestMemoryRegression:
    def test_simple_write_read(self, backend):
        """最基本写入再读出，序列化/反序列化一致性"""
        msgs = [{"role": "user", "content": "我是杭州人"}, {"role": "assistant", "content": "记住了"}]
        backend.save_memory("u1", "s1", msgs)
        loaded = backend.load_memory("u1", "s1")
        assert loaded == msgs

    def test_user_isolation(self, backend):
        """不同用户记忆严格隔离"""
        backend.save_memory("u1", "s1", [{"role": "user", "content": "A"}])
        backend.save_memory("u2", "s1", [{"role": "user", "content": "B"}])
        assert backend.load_memory("u1", "s1") == [{"role": "user", "content": "A"}]
        assert backend.load_memory("u2", "s1") == [{"role": "user", "content": "B"}]

    def test_memory_persistence_after_restart(self, docker_services, backend):
        """模拟服务重启后记忆不丢------重启容器再查"""
        msgs = [{"role": "user", "content": "重启前消息"}]
        backend.save_memory("u3", "s3", msgs)
        
        # 重启对应存储容器
        project_name = f"memtest_{os.getpid()}"
        client = docker.from_env()
        if isinstance(backend, RedisMemoryBackend):
            name_filter = "redis"
        else:
            name_filter = "postgres"
        for c in client.containers.list(filters={"name": project_name}):
            if name_filter in c.name:
                c.restart()
                break
        # 等待恢复
        time.sleep(3)
        # 重新连接后读取
        loaded = backend.load_memory("u3", "s3")
        assert loaded == msgs

踩坑实录：官方文档不会告诉你的两个坑

坑1：Docker 容器起来后端口监听未就绪，pytest 连接直接抛 ConnectionRefused。

表象：偶尔 CI 上刚 up 完马上跑用例就挂。原因是 docker-compose up -d 命令返回不代表服务可用，Redis 直到加载 RDB 完才 accept。解决是在 fixture 里用 redis-cli ping 和 pg_isready 做显式轮询，最多等 15 秒。千万别学有些文章只 time.sleep(2) 靠运气。

坑2：pytest-docker 插件的 session scope 导致数据库状态跨测试文件污染。

最初用了 pytest-docker 的 docker_ip 和 docker_services fixture，但插件强制 session scope，我两个 test 文件共用同一套 Redis，上一个文件清空了数据库，下一个文件还在用旧引用，出现诡异的 key 不存在错误。后来改用手动 module scope fixture，每个测试模块启动独立的 docker-compose project，状态完美隔离。这个坑浪费了我一个下午。

效果验证：从 3 小时到 5 分钟，还不用人盯着

优化前：每次版本发布前，QA 手动验证记忆存储回归，搭建环境 20 分钟，执行全用例（约 20 个场景）2.5 小时，经常忘记清理导致数据污染重来。

优化后：一行 pytest test_memory_regression.py，5 分钟跑完 16 个参数化场景，Redis/Postgres 全覆盖，含容器重启恢复测试。CI 上每次 MR 自动触发，发布再也没出现过"记忆消失"的 P0 单。回归测试从瓶颈变成了安全感来源。

# 你唯一需要跑的命令：
docker-compose -f docker-compose.yml up -d
pytest -v test_memory_regression.py

可以直接抄走的配置模板

把上面的 docker-compose.yml、conftest.py、测试文件拷进项目，改成自己的记忆存储接口实现，立刻拥有 AI 记忆回归能力。记得 pip install pytest redis psycopg2-binary docker。

---

#Python #Pytest #Docker #AI测试 #自动化测试

关于作者

一名给业务擦屁股擦出感情的实战派后端架构师，专注把"能跑就行"变成"不会翻车"。

GitHub: github.com/baofugege

Sponsor: github.com/sponsors/ba... --- 如果这篇方案帮你省下了 3 小时手动回归，请我喝杯咖啡

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