Pytest - 如何在每个测试函数后删除创建的数据

Question

我有一个使用FastAPI和SQLAlchemy的项目，并且我用Pytest来写API的单元测试。

在每个测试函数中，我会用SQLAlchemy在一些表里创建一些数据（比如用户表、帖子表、评论表等）。这些在每个测试函数中创建的数据在测试结束后仍然会留在表里，这样就会影响到其他测试函数。

举个例子，在第一个测试函数中我创建了3个帖子和2个用户，然后在第二个测试函数中，这3个帖子和2个用户依然在表里，这就导致我的测试结果不正确。

下面是我为pytest写的一个设置：

@pytest.fixture
def session(engine):
    Session = sessionmaker(bind=engine)
    session = Session()
    yield session
    session.rollback()  # Removes data created in each test method
    session.close()  # Close the session after each test

我用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

在最新的测试函数中，我本该得到404错误，但却得到了200，并且有5个帖子（是上一个测试函数的结果）

我该如何在每个测试函数结束后删除创建的数据呢？

Answer 1

一种解决办法是设置一个工具，它可以在每次测试之前清空数据库中的表格。这样你可以在类上像这样使用它：

import pytest
from sqlalchemy import text
from sqlalchemy.orm import Session

from models import user, post


@pytest.fixture()
def clean_db(session: Session):
    tables = [user.__tablename__, post.__tablename__]
    for table in tables:
        session.execute(text(f'TRUNCATE TABLE {table}'))
    session.commit()

然后在测试类中：

import pytest

@pytest.mark.usefixtures("clean_db", autouse=True)
class TestAllPosts(PostBaseTestCase):

...

Answer 2

首先，理想情况下，你的单元测试应该是相互独立的。也许你需要考虑重新设计一下它们，这样它们就可以随机执行。我知道这样做有时候会多花一些功夫……

我使用了 session.rollback() 来删除会话中创建的所有数据，但它并没有删除数据。

这是因为你已经提交了更改。如果你提交了更改，rollback 就不会影响这些更改。它只会影响当前正在进行的事务。

这里有两种可能的解决方案：

1. 不要提交。在你的测试中打开一个事务后，创建你的帖子和其他内容，然后不要提交它们，这样它们会随着你当前的设置一起回滚，不会影响数据库。你仍然可以读取这些更改，并确认你的测试确实有效。（想了解更多信息，可以查看隔离级别）

注意：你应该为你的设置指定 "function" 的作用域，这样每个测试都会进行回滚。

注意：记住，如果你在测试中使用 with session.begin() 上下文管理器，它会在退出 with 块时提交更改。如果你选择了这个解决方案，就应该避免这样做。只需使用 session.begin() 就足够了，用来打开一个事务。

2. 在你的测试中，删除你创建的对象，然后再次 提交 会话。

Answer 3

问题在于有多个会话。

一个会话是你测试时用的，另一个（或多个）会话是服务器在用的。

因为你使用了client.get，这就意味着你向服务器发送了请求，服务器会使用它自己的数据库会话。

1. 要解决这个问题，你可以在每次测试结束时清空所有表格：https://stackoverflow.com/a/25220958/5521670

@pytest.fixture
def session(engine):
    Session = sessionmaker(bind=engine)
    session = Session()
    yield session

    # Remove any data from database (even data not created by this session)
    with contextlib.closing(engine.connect()) as connection:
        transaction = connection.begin()
        connection.execute(f'TRUNCATE TABLE {",".join(table.name for table in reversed(Base.metadata.sorted_tables)} RESTART IDENTITY CASCADE;'))
        transaction.commit()

    session.rollback()  # Removes data created in each test method
    session.close()  # Close the session after each test

2. 另一个办法是让服务器使用你的测试会话（就像FastAPI文档中建议的那样）：https://fastapi.tiangolo.com/advanced/testing-database/

def override_get_db():
    try:
        db = TestingSessionLocal()
        yield db
    finally:
        db.close()


app.dependency_overrides[get_db] = override_get_db