python - 如何在 FastAPI 中的测试之间设置和拆除数据库？

Question

我已经按照 FastAPI documentation 设置了我的单元测试，但它仅涵盖数据库在测试中持久化的情况。
如果我想每次测试建立和拆除数据库怎么办？ (例如，下面的第二个测试会失败，因为第一个测试后数据库将不再为空)。
我目前正在通过拨打 create_all 来做这件事和 drop_all (在下面的代码中注释掉)在每次测试的开始和结束，但这显然不理想(如果测试失败，数据库永远不会被拆除，影响下一次测试的结果)。
我怎样才能正确地做到这一点？我应该在 override_get_db 周围创建某种 Pytest 固定装置吗？依赖？

from fastapi.testclient import TestClient
from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

from main import app, get_db
from database import Base

SQLALCHEMY_DATABASE_URL = "sqlite:///./test.db"

engine = create_engine(
    SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False}
)
TestingSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

# Base.metadata.create_all(bind=engine)

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

app.dependency_overrides[get_db] = override_get_db

client = TestClient(app)

def test_get_todos():
    # Base.metadata.create_all(bind=engine)

    # create
    response = client.post('/todos/', json={'text': 'some new todo'})
    data1 = response.json()
    response = client.post('/todos/', json={'text': 'some even newer todo'})
    data2 = response.json()

    assert data1['user_id'] == data2['user_id']

    response = client.get('/todos/')

    assert response.status_code == 200
    assert response.json() == [
        {'id': data1['id'], 'user_id': data1['user_id'], 'text': data1['text']},
        {'id': data2['id'], 'user_id': data2['user_id'], 'text': data2['text']}
    ]

    # Base.metadata.drop_all(bind=engine)

def test_get_empty_todos_list():
    # Base.metadata.create_all(bind=engine)

    response = client.get('/todos/')

    assert response.status_code == 200
    assert response.json() == []

    # Base.metadata.drop_all(bind=engine)

Answer 1

为了在测试失败后进行清理(并在测试前进行设置)，pytest 提供了 pytest.fixture .
在您的情况下，您希望在每次测试之前创建所有表，然后再次删除它们。这可以通过以下夹具来实现:

@pytest.fixture()
def test_db():
    Base.metadata.create_all(bind=engine)
    yield
    Base.metadata.drop_all(bind=engine)

然后在您的测试中使用它，如下所示:

def test_get_empty_todos_list(test_db):
    response = client.get('/todos/')

    assert response.status_code == 200
    assert response.json() == []

对于每个具有 test_db 的测试在其参数列表中 pytest 首先运行 Base.metadata.create_all(bind=engine) ，然后屈服于测试代码，然后确保 Base.metadata.drop_all(bind=engine)运行，即使测试失败。
完整代码:

import pytest
from fastapi.testclient import TestClient
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from main import app, get_db
from database import Base


SQLALCHEMY_DATABASE_URL = "sqlite:///./test.db"

engine = create_engine(
    SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False}
)
TestingSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)


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


@pytest.fixture()
def test_db():
    Base.metadata.create_all(bind=engine)
    yield
    Base.metadata.drop_all(bind=engine)

app.dependency_overrides[get_db] = override_get_db

client = TestClient(app)


def test_get_todos(test_db):
    response = client.post("/todos/", json={"text": "some new todo"})
    data1 = response.json()
    response = client.post("/todos/", json={"text": "some even newer todo"})
    data2 = response.json()

    assert data1["user_id"] == data2["user_id"]

    response = client.get("/todos/")

    assert response.status_code == 200
    assert response.json() == [
        {"id": data1["id"], "user_id": data1["user_id"], "text": data1["text"]},
        {"id": data2["id"], "user_id": data2["user_id"], "text": data2["text"]},
    ]


def test_get_empty_todos_list(test_db):
    response = client.get("/todos/")

    assert response.status_code == 200
    assert response.json() == []

随着应用程序的增长，为每个测试设置和拆除整个数据库可能会变慢。
一个解决方案是只设置数据库一次，然后从不实际向它提交任何内容。
这可以使用嵌套事务和回滚来实现:

import pytest
import sqlalchemy as sa
from fastapi.testclient import TestClient
from sqlalchemy.orm import sessionmaker

from database import Base
from main import app, get_db

SQLALCHEMY_DATABASE_URL = "sqlite:///./test.db"

engine = sa.create_engine(
    SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False}
)
TestingSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

# Set up the database once
Base.metadata.drop_all(bind=engine)
Base.metadata.create_all(bind=engine)


# These two event listeners are only needed for sqlite for proper
# SAVEPOINT / nested transaction support. Other databases like postgres
# don't need them. 
# From: https://docs.sqlalchemy.org/en/14/dialects/sqlite.html#serializable-isolation-savepoints-transactional-ddl
@sa.event.listens_for(engine, "connect")
def do_connect(dbapi_connection, connection_record):
    # disable pysqlite's emitting of the BEGIN statement entirely.
    # also stops it from emitting COMMIT before any DDL.
    dbapi_connection.isolation_level = None


@sa.event.listens_for(engine, "begin")
def do_begin(conn):
    # emit our own BEGIN
    conn.exec_driver_sql("BEGIN")


# This fixture is the main difference to before. It creates a nested
# transaction, recreates it when the application code calls session.commit
# and rolls it back at the end.
# Based on: https://docs.sqlalchemy.org/en/14/orm/session_transaction.html#joining-a-session-into-an-external-transaction-such-as-for-test-suites
@pytest.fixture()
def session():
    connection = engine.connect()
    transaction = connection.begin()
    session = TestingSessionLocal(bind=connection)

    # Begin a nested transaction (using SAVEPOINT).
    nested = connection.begin_nested()

    # If the application code calls session.commit, it will end the nested
    # transaction. Need to start a new one when that happens.
    @sa.event.listens_for(session, "after_transaction_end")
    def end_savepoint(session, transaction):
        nonlocal nested
        if not nested.is_active:
            nested = connection.begin_nested()

    yield session

    # Rollback the overall transaction, restoring the state before the test ran.
    session.close()
    transaction.rollback()
    connection.close()


# A fixture for the fastapi test client which depends on the
# previous session fixture. Instead of creating a new session in the
# dependency override as before, it uses the one provided by the
# session fixture.
@pytest.fixture()
def client(session):
    def override_get_db():
        yield session

    app.dependency_overrides[get_db] = override_get_db
    yield TestClient(app)
    del app.dependency_overrides[get_db]


def test_get_empty_todos_list(client):
    response = client.get("/todos/")

    assert response.status_code == 200
    assert response.json() == []

这里有两个装置( session 和 client )还有一个额外的优势:
如果测试只与 API 对话，那么您不需要记住显式添加 db 固定装置(但它仍然会被隐式调用)。
如果你想编写一个直接与数据库对话的测试，你也可以这样做:

def test_something(session):
    session.query(...)

或者两者兼而有之，例如，如果您想在 API 调用之前准备数据库状态:

def test_something_else(client, session):
    session.add(...)
    session.commit()
    client.get(...)

应用程序代码和测试代码都将看到数据库的相同状态。