如何在每次测试后将PostgreSQL数据库恢复至基准状态？

完全可行，你可以通过PostgreSQL的基准备份恢复或者模板数据库复制实现高效的状态重置，既解决提交后无法回滚的问题，又比重建数据库/手动写恢复SQL更高效易维护。以下是具体实现方案：

方案一：基准SQL备份+测试前恢复 #

适合数据量不大的场景，操作简单直接：

1. 会话级容器初始化（全局夹具） #

启动PostgreSQL容器后，先构建好仅含admin行的基准状态，然后导出备份文件：

import pytest
from testcontainers.postgres import PostgresContainer
from sqlalchemy import create_engine, text

@pytest.fixture(scope="session")
def postgres_container():
    with PostgresContainer("postgres:15") as container:
        # 连接默认数据库，初始化基准表和数据
        engine = create_engine(container.get_connection_url())
        with engine.connect() as conn:
            # 创建users表
            conn.execute(text("""
                CREATE TABLE users (
                    id SERIAL PRIMARY KEY,
                    username VARCHAR(50) UNIQUE NOT NULL
                )
            """))
            # 插入admin基准行
            conn.execute(text("INSERT INTO users (username) VALUES ('admin')"))
            conn.commit()
        
        # 导出基准状态到容器内的备份文件
        container.exec(f"pg_dump -U {container.username} -d {container.dbname} -f /tmp/base_backup.sql")
        yield container

2. 函数级状态重置（测试夹具） #

每个测试执行前，从备份文件恢复基准状态：

@pytest.fixture(scope="function")
def fresh_db(postgres_container):
    # 恢复基准备份
    postgres_container.exec(f"psql -U {postgres_container.username} -d {postgres_container.dbname} -f /tmp/base_backup.sql")
    # 返回SQLAlchemy连接供测试使用
    engine = create_engine(postgres_container.get_connection_url())
    with engine.connect() as conn:
        yield conn

3. 测试场景实现 #

每个测试直接使用fresh_db夹具，自动基于基准状态运行：

def test_only_admin(fresh_db):
    result = fresh_db.execute(text("SELECT username FROM users")).fetchall()
    assert len(result) == 1
    assert result[0][0] == "admin"

def test_admin_and_user_foo(fresh_db):
    # 添加测试数据
    fresh_db.execute(text("INSERT INTO users (username) VALUES ('user_foo')"))
    fresh_db.commit()
    # 验证状态
    usernames = [row[0] for row in fresh_db.execute(text("SELECT username FROM users")).fetchall()]
    assert len(usernames) == 2
    assert {"admin", "user_foo"} == set(usernames)

def test_admin_and_three_users(fresh_db):
    test_users = ["user_x", "user_y", "user_z"]
    for user in test_users:
        fresh_db.execute(text("INSERT INTO users (username) VALUES (:user)"), {"user": user})
    fresh_db.commit()
    # 验证状态
    usernames = [row[0] for row in fresh_db.execute(text("SELECT username FROM users")).fetchall()]
    assert len(usernames) == 4
    assert "admin" in usernames
    assert set(test_users).issubset(usernames)

方案二：模板数据库复制（更高效） #

如果你的测试涉及较多表或数据，推荐用PostgreSQL的模板数据库特性——它是文件级复制，速度远快于执行SQL备份：

1. 会话级模板初始化 #

创建基准数据库并标记为模板：

import pytest
import os
from testcontainers.postgres import PostgresContainer
from sqlalchemy import create_engine, text

@pytest.fixture(scope="session")
def postgres_container():
    with PostgresContainer("postgres:15") as container:
        # 创建基准数据库
        engine = create_engine(container.get_connection_url())
        with engine.connect() as conn:
            conn.execute(text("CREATE DATABASE base_db"))
            conn.commit()
        
        # 初始化基准数据库的表和数据
        base_engine = create_engine(container.get_connection_url(dbname="base_db"))
        with base_engine.connect() as conn:
            conn.execute(text("""
                CREATE TABLE users (
                    id SERIAL PRIMARY KEY,
                    username VARCHAR(50) UNIQUE NOT NULL
                )
            """))
            conn.execute(text("INSERT INTO users (username) VALUES ('admin')"))
            conn.commit()
        
        # 将基准数据库设置为模板
        container.exec(f"psql -U {container.username} -d postgres -c 'ALTER DATABASE base_db IS_TEMPLATE = TRUE;'")
        yield container

2. 函数级测试数据库创建 #

每个测试前从模板复制新数据库，测试后销毁：

@pytest.fixture(scope="function")
def test_db(postgres_container):
    # 生成唯一测试数据库名
    test_db_name = f"test_db_{hash(os.urandom(4))}"
    # 从模板复制新数据库
    postgres_container.exec(f"psql -U {postgres_container.username} -d postgres -c 'CREATE DATABASE {test_db_name} TEMPLATE base_db;'")
    # 返回测试数据库连接
    engine = create_engine(postgres_container.get_connection_url(dbname=test_db_name))
    with engine.connect() as conn:
        yield conn
    # 测试结束后清理数据库
    postgres_container.exec(f"psql -U {postgres_container.username} -d postgres -c 'DROP DATABASE {test_db_name};'")

3. 测试场景实现 #

和方案一的测试代码完全一致，只需把fresh_db换成test_db即可。

方案对比 #

• 备份恢复：实现简单，适合小型测试场景；恢复速度取决于SQL脚本大小。
• 模板复制：速度极快（毫秒级），适合表多、数据量大的场景；需要额外管理模板数据库。

这两种方案都能完美满足你的需求：每个测试都基于纯净的基准状态运行，测试后的变更不会影响其他测试，且无需手动维护恢复逻辑。

内容的提问来源于stack exchange，提问作者Vito De Tullio