如何在 FastAPI 测试之间建立和拆除数据库？

Question

我已经按照 FastAPI文档设置了单元测试，但它只涵盖了数据库在测试中保留的情况。

如果我想在每次测试时构建和拆除数据库怎么办？（例如，下面的第二个测试将失败，因为第一个测试后数据库将不再为空）。

我目前通过在每个测试的开始和结束时调用 and （在下面的代码中注释掉）来create_all做到这一点，但这显然不理想（如果测试失败，数据库将永远不会被拆除，影响下一个测试的结果drop_all测试）。

我怎样才能正确地做到这一点？我应该围绕依赖关系创建某种 Pytest 固定装置吗override_get_db？

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_dbpytest 首先运行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 对话，那么您不需要记住显式添加数据库固定装置（但仍会隐式调用它）。如果您想编写一个直接与数据库对话的测试，您也可以这样做：

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

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

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

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

Answer 2

这是完整的 FastAPI 测试环境的解决方案，包括数据库设置和拆卸。尽管事实上已经有一个公认的答案，但我还是想贡献我的想法。

配置测试环境时，您需要将这些装置包含在 conftest.py 文件中。测试包中包含的任何测试都可以自动访问其中定义的夹具。

a) 首先，进行导入。

请记住，您的导入路径可能与我的不同，因此也要仔细检查。

import pytest
from fastapi.testclient import TestClient

# Import the SQLAlchemy parts
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.declarative import declarative_base
from app.main import app
from app.database import get_db,Base

# Create the new database session

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

engine = create_engine(SQLALCHEMY_DATABASE_URL)

TestingSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

接下来，我们将使用 Pytest 装置，它们是在应用它们的每个测试函数之前运行的函数。

b). 会议赛程

@pytest.fixture()
def session():

    Base.metadata.drop_all(bind=engine)
    Base.metadata.create_all(bind=engine)

    db = TestingSessionLocal()

    try:
        yield db
    finally:
        db.close()

上面的会话固定装置确保每次运行测试时，我们都会连接到测试数据库，创建表，然后在测试完成后删除这些表。

c) 客户端固定装置

@pytest.fixture()
def client(session):

    # Dependency override

    def override_get_db():
        try:

            yield session
        finally:
            session.close()

    app.dependency_overrides[get_db] = override_get_db

    yield TestClient(app)

上面的装置将我们连接到新的测试数据库，并覆盖主应用程序建立的初始数据库连接。该客户端固定装置需要会话固定装置才能运行。

之后，您可以使用如图所示的装置，而无需导入任何内容，如下所示。

def test_index(client):
    res = client.get("/")
    assert res.status_code == 200

您的完整 conftest.py 文件现在应如下所示：

import pytest
from fastapi.testclient import TestClient

# Import the SQLAlchemy parts
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.declarative import declarative_base
from app.main import app
from app.database import get_db, Base

# Create the new database session

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

engine = create_engine(SQLALCHEMY_DATABASE_URL)

TestingSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)


@pytest.fixture()
def session():

    # Create the database

    Base.metadata.drop_all(bind=engine)
    Base.metadata.create_all(bind=engine)

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


@pytest.fixture()
def client(session):

    # Dependency override

    def override_get_db():
        try:
            yield session
        finally:
            session.close()

    app.dependency_overrides[get_db] = override_get_db

    yield TestClient(app)