SQLAlchemy+pytest异步测试中两种数据库回滚方案的对比与验证

我正在开发一个基于FastAPI和异步SQLAlchemy的项目，已经为项目编写了pytest测试，并且成功实现了每次测试后的数据库回滚。目前我找到了两种不同的实现方式，但不确定二者之间的差异，想知道这两种方法是否都正确，或是其中某一种存在潜在问题。

conftest.py 代码实现 #

# pyproject.toml
#
# pytest = "^8.3.2"
# pytest-asyncio = "==0.21.2"
# #pytest-dependency = "^0.6.0"
# pytest-order = "^1.2.1"
#
# [tool.pytest.ini_options]
# addopts = "-s"
# asyncio_mode = "auto"

import asyncio
from urllib.parse import urlparse

import pytest
from sqlalchemy import NullPool
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession

from config import settings
from depends.db import async_session as api_async_session
from main import app

url = urlparse(settings.db)._replace(scheme="postgresql+asyncpg").geturl()


@pytest.fixture(scope="session")
def event_loop(request):
    loop = asyncio.get_event_loop_policy().new_event_loop()
    yield loop
    loop.close()


@pytest.fixture
async def async_session():
    async_db = create_async_engine(url, echo=False, poolclass=NullPool)
    async with async_db.connect() as connection:
        async with connection.begin() as transaction:
            async with AsyncSession(bind=connection) as s:
                app.dependency_overrides[api_async_session] = lambda: s
                yield s
            await transaction.rollback()

# Instead of using a connection pool, bind a specific connection to the event loop. 
# If you don't set an event loop policy or using pytest-asyncio 0.23, 
# each test will start a new event loop, causing asyncpg triggering exceptions.
@pytest.fixture
async def async_session2():
    async_db = create_async_engine(url, echo=False, poolclass=NullPool)
    async with async_db.connect() as connection:
        transaction = await connection.begin()
        async with AsyncSession(bind=connection, join_transaction_mode="create_savepoint") as s:
            app.dependency_overrides[api_async_session] = lambda: s
            yield s
        await transaction.rollback()

两种实现的核心差异与正确性分析 #

1. 事务管理方式的区别 #

• async_session 实现：
  采用async with connection.begin() as transaction上下文管理器模式，进入上下文时自动开启事务，正常退出时会自动提交事务，但代码在yield后手动调用await transaction.rollback()，覆盖了自动提交逻辑，确保无论测试成功或失败，最终都会回滚事务。同时，如果测试过程中抛出异常，上下文管理器会自动触发事务回滚，无需额外处理。

• async_session2 实现：
  手动调用transaction = await connection.begin()开启事务，没有使用上下文管理器包裹事务逻辑。事务的提交/回滚完全依赖手动代码控制，最终通过await transaction.rollback()完成回滚。若在yield之前（比如创建AsyncSession阶段）抛出异常，事务不会自动回滚，但由于连接被async with async_db.connect()包裹，连接关闭时数据库会自动回滚未提交的事务，因此不会留下脏数据。

2. 会话与事务的绑定模式差异 #

• async_session 中的AsyncSession：
  未指定join_transaction_mode，默认使用"auto"模式。当连接已处于事务中时，会话会直接加入该事务，会话的commit()操作不会真正提交事务，仅会将变更刷新到当前连接的事务上下文里，最终由连接级事务统一处理。

• async_session2 中的AsyncSession：
  指定join_transaction_mode="create_savepoint"，会话会在当前连接的事务内创建一个保存点（Savepoint）。此时会话的commit()仅会提交保存点内的变更（将数据写入连接事务），而会话的rollback()只会回滚到该保存点，不会影响连接级的整体事务。不过在当前测试场景下，最终都会回滚连接级事务，因此两种模式的最终效果一致。

3. 潜在问题对比 #

• async_session：
  符合Python上下文管理器的最佳实践，异常处理更安全，几乎没有潜在风险。唯一需要注意的是，若测试中手动调用了连接的rollback()，可能会导致事务提前结束，后续数据库操作报错，但这种场景在测试中很少出现。

• async_session2：
  手动管理事务的方式灵活性更高（适合需要在会话内做部分回滚的场景），但相较于上下文管理器模式，异常处理的容错性稍弱——如果yield前出现异常，事务无法自动回滚，不过依赖连接关闭时的数据库自动回滚机制，也不会产生持久化的脏数据。

总结 #

两种实现方式都能正确实现测试后的数据库回滚需求：

• 若追求代码简洁、符合最佳实践，优先选择async_session的实现；
• 若测试中需要会话级的部分回滚操作，async_session2的create_savepoint模式会更适用。

内容的提问来源于stack exchange，提问作者ACE Fly