如何在 FastAPI 的测试之间建立和拆除数据库？
How to set up and tear down a database between tests in FastAPI?
我已经按照 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)
为了在测试失败后进行清理（并在测试前进行设置），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() == []
对于每个在其参数列表 pytest 中具有 test_db 的测试，首先 运行s 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 fixture（但仍会隐式调用）。
如果你想编写一个直接与数据库对话的测试，你也可以这样做：
def test_something(session):
session.query(...)
或两者，例如，如果您想要在 API 调用之前准备数据库状态：
def test_something_else(client, session):
session.add(...)
session.commit()
client.get(...)
应用程序代码和测试代码都将看到相同的数据库状态。
您还可以在每次测试后运行分类表格运行。这会清除其中的所有数据，而不会实际删除模式，因此它不会像 Base.metadata.drop_all(bind=engine):
那样慢
import pytest
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, declarative_base
from contextlib import contextmanager
engine = create_engine('postgresql://...')
Session = sessionmaker(bind=engine)
Base = declarative_base()
@contextmanager
def session_scope():
"""Provide a transactional scope around a series of operations."""
session = Session()
try:
yield session
session.commit()
except:
session.rollback()
raise
finally:
session.close()
def clear_tables():
with session_scope() as conn:
for table in Base.metadata.sorted_tables:
conn.execute(
f"TRUNCATE {table.name} RESTART IDENTITY CASCADE;"
)
conn.commit()
@pytest.fixture
def test_db_session():
yield engine
engine.dispose()
clear_tables()
def test_some_feature(test_db_session):
test_db_session.query(...)
(...)
这是完整的 FastAPI 测试环境的解决方案，包括数据库设置和拆卸。尽管已经有一个公认的答案，但我还是想贡献我的想法。
配置测试环境时，您需要将这些装置包含在 conftest.py 文件中。其中定义的夹具将自动可供测试包中包含的任何测试访问。
a) 首先，进行导入。
Remember that your imports path may differ from mine, so double-check that as well.
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). Session 夹具
@pytest.fixture()
def session():
Base.metadata.drop_all(bind=engine)
Base.metadata.create_all(bind=engine)
db = TestingSessionLocal()
try:
yield db
finally:
db.close()
上面的 session fixture 确保每次测试 运行，我们连接到测试数据库，创建表，然后在测试完成后删除表。
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)
上面的装置将我们连接到新的测试数据库，并覆盖了主应用程序建立的初始数据库连接。此客户端装置需要 session 装置才能运行。
之后，您可以使用如图所示的灯具，而无需导入任何东西，如下所示。
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)
我已经按照 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)
为了在测试失败后进行清理（并在测试前进行设置），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() == []
对于每个在其参数列表 pytest 中具有 test_db 的测试，首先 运行s 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 fixture（但仍会隐式调用）。 如果你想编写一个直接与数据库对话的测试，你也可以这样做：
def test_something(session):
session.query(...)
或两者，例如，如果您想要在 API 调用之前准备数据库状态：
def test_something_else(client, session):
session.add(...)
session.commit()
client.get(...)
应用程序代码和测试代码都将看到相同的数据库状态。
您还可以在每次测试后运行分类表格运行。这会清除其中的所有数据，而不会实际删除模式，因此它不会像 Base.metadata.drop_all(bind=engine):
那样慢import pytest
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, declarative_base
from contextlib import contextmanager
engine = create_engine('postgresql://...')
Session = sessionmaker(bind=engine)
Base = declarative_base()
@contextmanager
def session_scope():
"""Provide a transactional scope around a series of operations."""
session = Session()
try:
yield session
session.commit()
except:
session.rollback()
raise
finally:
session.close()
def clear_tables():
with session_scope() as conn:
for table in Base.metadata.sorted_tables:
conn.execute(
f"TRUNCATE {table.name} RESTART IDENTITY CASCADE;"
)
conn.commit()
@pytest.fixture
def test_db_session():
yield engine
engine.dispose()
clear_tables()
def test_some_feature(test_db_session):
test_db_session.query(...)
(...)
这是完整的 FastAPI 测试环境的解决方案，包括数据库设置和拆卸。尽管已经有一个公认的答案，但我还是想贡献我的想法。
配置测试环境时，您需要将这些装置包含在 conftest.py 文件中。其中定义的夹具将自动可供测试包中包含的任何测试访问。
a) 首先，进行导入。
Remember that your imports path may differ from mine, so double-check that as well.
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). Session 夹具
@pytest.fixture()
def session():
Base.metadata.drop_all(bind=engine)
Base.metadata.create_all(bind=engine)
db = TestingSessionLocal()
try:
yield db
finally:
db.close()
上面的 session fixture 确保每次测试 运行，我们连接到测试数据库，创建表，然后在测试完成后删除表。
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)
上面的装置将我们连接到新的测试数据库，并覆盖了主应用程序建立的初始数据库连接。此客户端装置需要 session 装置才能运行。
之后，您可以使用如图所示的灯具，而无需导入任何东西，如下所示。
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)