Pytest与SQLAlchemy测试间事务回滚失败问题求助

问题概述 #

采用经典/命令式映射的领域模型，使用SQLAlchemy编写测试时无法实现测试间事务回滚，前一个测试插入的数据会遗留到下一个测试中，导致断言失败。已确认：

• SQLite环境下问题源于SQLite与SQLAlchemy嵌套事务（SAVEPOINT）的兼容性问题
• SQL Server环境下事务可回滚，但自增主键序列不会随事务回滚重置

代码重现 #

model.py #

from dataclasses import dataclass

@dataclass(frozen=False)
class Dummy:
    _id: int
    value: int

orm.py #

from sqlalchemy import Table, Column, Integer
from sqlalchemy.orm import registry

mapper_registry = registry()

dummy = Table(
    "dummy",
    mapper_registry.metadata,
    Column("_id", Integer, primary_key=True, autoincrement=True),
    Column("value", Integer, nullable=False)
)

def start_mappers():
    mapper_registry.map_imperatively(Dummy, dummy)

test_orm.py #

def test_inmem_dummy_mapped(session):
    dummy = Dummy(1, 1)
    with session:
       session.add(dummy)
       session.add(Dummy(2,4))
       session.commit()

    results = session.query(Dummy).all()
    assert results == [Dummy(1, 1), Dummy(2,4)]

def test_inmem_nested_transaction_A(session):
    dummy = Dummy(3, 1)
    with session:
       session.add(dummy)
       session.add(Dummy(4,6))
       session.commit()

    results = session.query(Dummy).all()
    assert results == [Dummy(3, 1), Dummy(4,6)]

conftest.py（原版本） #

import pytest
from sqlalchemy import create_engine
from sqlalchemy.orm import Session
from sqlalchemy import event
from orm import mapper_registry, start_mappers

@pytest.fixture(scope="session")
def engine():
    engine = create_engine("sqlite:///:memory:", echo=True, future=True)
    return engine

@pytest.fixture(scope="session")
def connection(request, engine):
    connection = engine.connect()
    def teardown():
        connection.close()
    request.addfinalizer(teardown)
    return connection

@pytest.fixture(scope="session", autouse=True)
def init_db_orm_mappings(connection, request):
    mapper_registry.metadata.bind = connection
    mapper_registry.metadata.create_all()
    start_mappers()
    def teardown():
        mapper_registry.metadata.drop_all()

    request.addfinalizer(teardown)

# 原session fixture版本1
# @pytest.fixture(scope="function", autouse = True)
# def session(connection, request, engine):
#     connection.commit()
#     transaction = connection.begin()
#     session = Session(bind=connection)
#     session.begin_nested()
#     yield session
#     transaction.rollback()

# 原session fixture版本2
# @pytest.fixture(scope="function", autouse = True)
# def session(connection, request, engine):
#     transaction = connection.begin()
#     session = Session(bind=connection)
#     session.begin_nested()

#     @event.listens_for(session, "after_transaction_end")
#     def restart_savepoint(db_session, transaction):
#         if transaction.nested and not transaction._parent.nested:
#             session.expire_all()
#             session.begin_nested()

#     def teardown():
#         session.close()
#         transaction.rollback()
#         session.rollback()

#     request.addfinalizer(teardown)
#     return session

解决方案 #

一、SQLite环境修复事务回滚 #

1. 集成SQLite兼容性补丁

修改conftest.py中的engine fixture，添加官方提供的事件监听代码，解决SAVEPOINT兼容性问题：

@pytest.fixture(scope="session")
def engine():
    engine = create_engine("sqlite:///:memory:", echo=True, future=True)
    
    # 修复SQLite嵌套事务兼容性
    @event.listens_for(engine, "connect")
    def do_connect(dbapi_connection, connection_record):
        # 禁用pysqlite自动发送BEGIN语句，阻止DDL前自动COMMIT
        dbapi_connection.isolation_level = None

    @event.listens_for(engine, "begin")
    def do_begin(conn):
        # 手动发送BEGIN语句
        conn.exec_driver_sql("BEGIN")
    
    return engine

2. 修正Session Fixture

使用以下版本的session fixture，确保嵌套事务正确工作：

@pytest.fixture(scope="function")
def session(connection, request):
    # 开启顶层事务
    transaction = connection.begin()
    # 创建绑定到当前连接的Session
    session = Session(bind=connection)
    
    # 开启嵌套事务（保存点）
    session.begin_nested()

    # 事务结束后自动重启嵌套事务，支持测试内多次commit
    @event.listens_for(session, "after_transaction_end")
    def restart_savepoint(db_session, tx):
        if tx.nested and not tx._parent.nested:
            db_session.expire_all()
            db_session.begin_nested()

    def teardown():
        # 回滚顶层事务，清除所有测试数据
        session.close()
        transaction.rollback()

    request.addfinalizer(teardown)
    return session

3. 调整测试代码（可选）

测试中避免使用with session:上下文（会自动提交顶层事务），直接使用session操作：

def test_inmem_dummy_mapped(session):
    dummy = Dummy(1, 1)
    session.add(dummy)
    session.add(Dummy(2,4))
    # 提交嵌套事务（保存点），不影响顶层事务
    session.commit()

    results = session.query(Dummy).all()
    assert results == [Dummy(1, 1), Dummy(2,4)]

二、SQL Server环境重置主键序列 #

SQL Server的自增IDENTITY序列不会随事务回滚重置，需手动重置：
修改session fixture，在测试前后重置表的主键序列：

@pytest.fixture(scope="function")
def session(connection, request):
    transaction = connection.begin()
    session = Session(bind=connection)
    
    # 测试开始前重置Dummy表的IDENTITY种子
    connection.exec_driver_sql("DBCC CHECKIDENT ('dummy', RESEED, 0);")
    
    session.begin_nested()

    @event.listens_for(session, "after_transaction_end")
    def restart_savepoint(db_session, tx):
        if tx.nested and not tx._parent.nested:
            db_session.expire_all()
            db_session.begin_nested()

    def teardown():
        session.close()
        transaction.rollback()
        # 也可选择在测试结束后重置
        # connection.exec_driver_sql("DBCC CHECKIDENT ('dummy', RESEED, 0);")

    request.addfinalizer(teardown)
    return session

内容的提问来源于stack exchange，提问作者Gary Frewin