Pytest技巧：如何在每个测试函数后清理创建的数据

问题背景 #

在FastAPI + SQLAlchemy项目中用Pytest编写API单元测试时，每个测试函数创建的用户、帖子、评论等数据会残留到下一个测试，导致预期结果错误。比如第一个测试创建5篇帖子后，第二个本该返回404的测试却返回200并携带之前的帖子。

当前使用的Pytest会话夹具如下：

@pytest.fixture
def session(engine):
    Session = sessionmaker(bind=engine)
    session = Session()
    yield session
    session.rollback()  # 尝试回滚清理测试数据
    session.close()

但session.rollback()并未生效，数据依然残留。

测试函数示例：

class TestAllPosts(PostBaseTestCase):

    def create_logged_in_user(self, db):
        user = self.create_user(db)
        return user.generate_tokens()["access"]

    def test_can_api_return_all_posts_without_query_parameters(self, client, session):
        posts_count = 5
        user_token = self.create_logged_in_user(session)
        for i in range(posts_count):
            self.create_post(session)

        response = client.get(url, headers={"Authorization": f"Bearer {user_token}"})
        assert response.status_code == 200
        json_response = response.json()
        assert len(json_response) == posts_count

    def test_can_api_detect_there_is_no_post(self, client, session):
        user_token = self.create_logged_in_user(session)
        response = client.get(url, headers={"Authorization": f"Bearer {user_token}"})
        assert response.status_code == 404

解决方案 #

1. 排查数据提交行为 #

session.rollback()无效的核心原因大概率是你在创建数据时主动提交了事务（比如调用了session.commit()）。如果数据已经提交到数据库，单纯的会话回滚无法撤销已提交的变更。

修改数据创建逻辑：

• 检查create_user、create_post等方法，移除其中的session.commit()调用
• 若需要获取自动生成的ID等数据，改用session.flush()替代提交，保持事务未闭合状态

2. 用连接级事务包裹测试 #

如果API内部会触发提交操作，无法避免事务提交，可以改用基于数据库连接的事务夹具，为每个测试创建独立的事务环境：

@pytest.fixture(autouse=True)
def db_session(engine):
    # 创建独立连接并开启事务
    connection = engine.connect()
    transaction = connection.begin()
    # 绑定当前连接到会话
    Session = sessionmaker(bind=connection)
    session = Session()

    yield session

    # 回滚会话和连接级事务，清除所有测试数据
    session.rollback()
    transaction.rollback()
    connection.close()

这个夹具会自动作用于所有测试，确保每个测试的变更都在独立事务中，结束后彻底回滚。

3. 测试专用数据库（可选） #

如果上述方法仍有问题，可以为测试单独配置数据库，在测试前后重建表结构：

@pytest.fixture(scope="session")
def test_engine():
    # 使用内存SQLite或独立测试数据库
    return create_engine("sqlite:///:memory:")

@pytest.fixture(scope="session")
def setup_tables(test_engine):
    from your_app.models import Base
    # 创建测试表
    Base.metadata.create_all(test_engine)
    yield
    # 测试结束后销毁表
    Base.metadata.drop_all(test_engine)

@pytest.fixture
def session(test_engine, setup_tables):
    connection = test_engine.connect()
    transaction = connection.begin()
    Session = sessionmaker(bind=connection)
    session = Session()
    yield session
    session.rollback()
    transaction.rollback()
    connection.close()

这种方式隔离性最强，但测试执行速度会稍慢，适合对环境纯净度要求高的场景。

关键注意事项 #

• 确保测试中使用的session和API实际依赖的是同一个会话/连接，否则API可能读取到未回滚的数据
• 避免在测试代码中直接调用数据库提交操作，所有变更都应放在事务中统一处理

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