如何在Flask-SQLAlchemy中实现测试后回滚主代码数据库会话

背景 #

我有一个大规模测试套件需要优化，最大瓶颈是每次测试都会删除并重建数据库表。如果整个测试会话仅执行一次删表与建表操作，测试执行时间可缩短至原来的1/4，因此切换为每次测试后回滚事务至关重要。

当前每个测试都使用以下夹具来删表、建表并添加测试数据：

from xxx.main import db

@pytest.fixture(scope="session")
def app():
    from xxx.main import application
    return TestApp(application)

# 每个测试都会使用这个夹具
@pytest.fixture(scope="function")
def drop_and_create(app):
    engine = create_engine("url_to_postgres")
    m = MetaData()
    m.reflect(engine)
    m.drop_all(engine)

    db.create_all()
    db.session.commit()
    # 下方添加测试数据

应用使用flask-sqlalchemy定义模型，与postgresql数据库交互：

# 模型定义
from xxx import db

class ModelName(db.Model):
    ...

db变量的定义如下：

import connexion
from flask_migrate import Migrate
app = connexion.App(__name__)
application = app.app
db = create_db(application)
migrate = Migrate(app=application, db=db)
metadata = db.metadata

def create_db(app):
    # 部分配置示例
    app.config["SQLALCHEMY_DATABASE_URI"] = ...
    app.config["SQLALCHEMY_ENGINE_OPTIONS"] = {
        "pool_recycle": app.config["SA_POOL_RECYCLE"],
        "pool_pre_ping": True,
    }
    db = SQLAlchemy(
        app,
        session_options={
            "autoflush": True,
            "autocommit": False,
            "expire_on_commit": False,
        },
    )
    db.init_app(app)
    return db

问题描述 #

我无法回滚主代码（端点处理函数）中产生的数据库变更：测试代码本身的db事务可以回滚，但端点处理函数中的事务无法回滚。

测试代码示例：

from xxx.main import db
def test_smth(app):
    response = cms_app.post(
        params=json.dumps(..),
        url="/register-user",                # 端点处理函数中的事务不会被回滚
        headers=headers,                     
        expect_errors=True,
        content_type="application/json",
    )
   
    db.add(something) 
    db.session.commit() # 该操作会被回滚

端点处理函数示例：

from xxx.main import db
# 端点定义
@blueprint.route("/register-user", methods=["POST"])
def registerUser():
    db.add(..)
    db.session.commit()   # 这些事务在测试结束后不会被回滚

已尝试方案 #

我试过一种接近需求的方案：它能回滚测试代码中的所有事务，但无法回滚主代码中的事务。测试速度变快，但部分测试失败，因为数据会持久化到数据库，测试结束后用pgAdmin查看仍有数据残留。

尝试的夹具代码：

# 每个测试都添加这个夹具
@pytest.fixture(scope="function")
def db_session(app):
    """
    创建嵌套事务，当应用代码调用session.commit时重建事务，测试结束后回滚。
    """
    connection = db.engine.connect()
    transaction = connection.begin()
    db.session = db.create_scoped_session(options={'bind': connection})
 
    # 开启嵌套事务（使用SAVEPOINT）
    nested = connection.begin_nested()
 
    # 如果应用代码调用session.commit，会结束嵌套事务，此时需要重新开启
    @listens_for(db.session, "after_transaction_end")
    def end_savepoint(session, transaction):
        nonlocal nested
        if not nested.is_active:
            nested = connection.begin_nested()
 
    yield
 
    # 回滚整体事务，恢复测试前的状态
    db.session.close()
    transaction.rollback()
    connection.close()

解决方案 #

问题核心是端点处理函数使用的db.session未被正确绑定到测试的事务连接上，导致端点的commit操作直接提交到数据库。以下是修复后的夹具方案：

1. 会话级数据库初始化夹具（仅执行一次） #

替换原来的drop_and_create，在整个测试会话开始时创建表并添加测试数据，结束时清理：

@pytest.fixture(scope="session")
def setup_database(app):
    # 创建所有数据库表
    db.create_all()
    # 初始化测试数据（原drop_and_create中的数据添加逻辑）
    # 示例：db.add(TestUser(name="test")); db.session.commit()
    db.session.commit()
    yield
    # 测试会话结束后删除所有表
    db.drop_all()

2. 函数级事务会话夹具（每个测试执行） #

确保所有代码（包括端点）使用的db.session都绑定到同一个带事务的连接，并拦截commit操作：

from sqlalchemy import event
from sqlalchemy.orm import scoped_session

@pytest.fixture(scope="function")
def transactional_session(setup_database):
    # 获取数据库连接并开启顶层事务
    connection = db.engine.connect()
    transaction = connection.begin()

    # 创建绑定到当前连接的scoped session
    session_factory = db.create_session(options={"bind": connection})
    scoped_session_obj = scoped_session(session_factory)

    # 替换全局的db.session，确保端点代码也使用该会话
    original_session = db.session
    db.session = scoped_session_obj

    # 开启嵌套事务（SAVEPOINT）
    nested = connection.begin_nested()

    # 拦截会话的commit操作：实际只回滚嵌套事务并重新开启，不真正提交到数据库
    @event.listens_for(scoped_session_obj, "after_transaction_end")
    def handle_after_transaction_end(session, tx):
        nonlocal nested
        if not nested.is_active:
            nested = connection.begin_nested()

    yield scoped_session_obj

    # 测试结束后清理：回滚顶层事务、关闭连接、恢复原有会话
    scoped_session_obj.remove()
    transaction.rollback()
    connection.close()
    db.session = original_session

3. 修改测试用例 #

让所有需要数据库操作的测试依赖transactional_session夹具：

def test_smth(app, transactional_session):
    # 测试逻辑不变，端点的事务会被自动回滚
    response = cms_app.post(...)
    # ...

原理说明 #

• 顶层事务：在测试开始时开启，测试结束后整体回滚，确保所有测试变更都被撤销。
• 嵌套事务（SAVEPOINT）：拦截代码中的commit操作，将其转化为嵌套事务的回滚与重建，避免真正提交到数据库。
• 全局会话替换：确保端点处理函数使用的db.session与测试代码使用的是同一个事务性会话，所有变更都在顶层事务的范围内。

内容的提问来源于stack exchange，提问作者S. Kostadinov