12
FastAPI Tips & Tricks: Testing a Database
If you haven't heard of it yet, FastAPI is a micro-framewok that allows developers to make full use of modern Python."Micro" here, means that rather than trying to cover every use case, it focuses on doing a single thing extremely well: giving you the tools to build a fast (duh) API. Everything else is up to you which allows you to tailor you architecture to perfectly fit your needs.
However it does mean that you have to set up yourself some things that usually comes out of the box in other frameworks. For example, if you want to use a relational database,
you'll have to choose and install an ORM. SQLAlchemy is the one documented by Fast API. There is a guide to help you set it up, and a tutorial wich gives some indications on how to go about testing it. But if like me, you come from Django, you might still struggle to configure everything so that each test works in isolation, and leaves the database as it was before running.
That's what we'll be covering in this article !
The code presented in this article is fully available on Github. Simply clone the repo and call docker-compose up to launch everything.
First things first, we need something to test ! We'll create a small project with an endpoint to create an "Item" with a title and a description and store it in the database. I personnaly prefer to use docker-compose to run a Fastapi image as well as a Postgres database next to it. But it's not mandatory.
Let's get to it then !
SQLAlchemy works by creating a declarative base, which allows us to describe our Database tables as Python classes that inherits it. Let's configure it in a file called database.py:
# database.py
from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
# We're using postgres but you could use
# any other engine supported by SQlAlchemy
SQLALCHEMY_DATABASE_URL = "postgresql://test-fastapi:password@db/db"
engine = create_engine(
SQLALCHEMY_DATABASE_URL
)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()Now we can create our Item model by inherting Base
For our Item model, we simply create a file called models.py and declare it using the Base we've juste configured :
# models.py
from sqlalchemy import Column
from sqlalchemy import Integer
from sqlalchemy import String
from .database import Base
class Item(Base):
__tablename__ = "items"
# Don't forget to set an index !
id = Column(Integer, primary_key=True, index=True)
title = Column(String, index=True)
description = Column(String, index=True)We also need to define the corresponding Pydantic schemas. FastAPI uses them to perform validation and serailization :
# schemas.py
from typing import Optional
from pydantic import BaseModel
class ItemBase(BaseModel):
title: "str"
description: "Optional[str] = None"
class ItemCreate(ItemBase):
pass
class Item(ItemBase):
id: int
class Config:
orm_mode = TrueAnd finally some CRUD utils for handling our Item instances :
# crud.py
from sqlalchemy import select
from sqlalchemy.orm import Session
from . import schemas
from .models import Item
def get_items(db: Session):
items = select(Item)
return db.execute(items).scalars().all()
def create_item(db: Session, item: schemas.ItemCreate):
db_item = Item(**item.dict())
db.add(db_item)
db.commit()
db.refresh(db_item)
return db_itemNow that we're done with everything database related, we can create the endpoint that we'll be testing !
Create a main.py file and add the following lines to it :
# main.py
from typing import List
from fastapi import Depends
from fastapi import FastAPI
from sqlalchemy.orm import Session
from . import crud
from . import models
from . import schemas
from .database import engine
from .database import SessionLocal
app = FastAPI()
# Here we create all the tables directly in the app
# in a real life situation this would be handle by a migratin tool
# Like alembic
models.Base.metadata.create_all(bind=engine)
# Dependency
def get_db():
db = SessionLocal()
try:
yield db
finally:
db.close()
@app.post("/items/", response_model=schemas.Item)
def create_item(item: schemas.ItemCreate, db: Session = Depends(get_db)):
return crud.create_item(db, item)Notice that we pass the SQLAlchemy Session as a dependency to our endpoint. This is important to note as it'll help us easily test our endpoint as we'll see in the next section.
We can check that our endpoint works thanks to the Swagger available on the /docs endpoint. Another awesome feature of FastAPI !
That's great! But our goal here is to test the endpoint automatically with some unit tests. We need to make sure that our tests don't affect the main database. We also want them to be deterministic and reusable, let's see how to do exactly that !
First let's try the method described in the official documentation. The idea here is to
leverage FastAPI Dependency system. Since our endpoint receives its session by dependency injection, we can use Dependency overrides to replace it with a session pointing to a test Database.
To to that, create a file called test_database.py to write our tests, and add the following code to it :
# test_database.py
SQLALCHEMY_DATABASE_URL = "postgresql://test-fastapi:password@db/test-fastapi-test"
engine = create_engine(SQLALCHEMY_DATABASE_URL)
TestingSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
def override_get_db():
try:
db = TestingSessionLocal()
yield db
finally:
db.close()
app.dependency_overrides[get_db] = override_get_dbThis will replace all instances of the get_db dependency with the override_get_db which return a Session connected to our test database. This
does one aspect of what we're looking for: we can run ou tests without affecting our main database.
Let's write a test and run it with pytest to see if that works as expected.
def test_post_items():
# We grab another session to check
# if the items are created
db = override_get_db()
client = TestClient(app)
client.post("/items/", json={"title": "Item 1"})
client.post("/items/", json={"title": "Item 2"})
items = crud.get_items(db)
assert len(items) == 2
It works perfectly ! Or does it ? We also want our tests to be deterministics, so if we run pytest again we should get the same result...
But we dont ! The items created the first time we ran the test are still in the database, so now we have
4 of them instead of 2. Now, we could delete the items manually. However, this is not something we want to do, because this will
become really tedious to maintain once we have more tests that validate more complex behaviors.
Luckily for use there IS a better way !
To solve our problem we'll make use of SQL Transactions. Essentially
transactions are a way to keep a set SQL instructions private to other database connections, until they are commited. And, if the changes
should not be commited, they can be rolled back ! Let's add try to add that to our dependency override. Disclaimer the following block of code doesn't actually works, we'll see why in a minute :
# The following doesn't work
# changes won't be rolled back !
def override_get_db():
try:
db = TestingSessionLocal()
db.begin()
yield db
finally:
db.rollback()
db.close()
app.dependency_overrides[get_db] = override_get_dbIn a FastAPI dependency, everything that comes after the yield instuction is executed after the path operation has run. So that's a great place to
try to roll back the changes made during the test. This looks good, but as I said, it doesn't work yet
That's because a call to commit always commit the outermost transaction, so we can't nest transactions thath way. This is a conscious choice from the SQLAlchemy developers to
avoid creating confusing and hard to predict behaviors.
Luckily for us they have provided an escape hatch designed to be used test suites. Indeed, a Session can be bound to an existing transaction, by opening it this way :
def override_get_db():
connection = engine.connect()
# begin a non-ORM transaction
transaction = connection.begin()
# bind an individual Session to the connection
db = Session(bind=connection)
# db = Session(engine)
yield db
db.rollback()
connection.close()And this time it will work !
Now we can run our test as many times as we want, and it will always pass. But we're not quite done yet. This works for data created by the endpoint,
but what if we wanted to seed the database during our test, say for testing an endpoint listing the available items ? The dependency override won't work in this case, because it only deals with
the dependency injection system.
Let's quickly implement the example I was just talking about: a READ endpoint that lists the existing items in the database :
@app.get("/items/", response_model=List[schemas.Item])
def read_items(db: Session = Depends(get_db)):
return crud.get_items(db)In order to test it we'd want to create a bunch of items before calling the endpoint, while having the database revert back to its original state after the test. The solution is simple : define our test database session in a fixture as well :
# conftest.py
@pytest.fixture(scope="session")
def db_engine():
engine = create_engine(SQLALCHEMY_DATABASE_URL)
if not database_exists:
create_database(engine.url)
Base.metadata.create_all(bind=engine)
yield engine
@pytest.fixture(scope="function")
def db(db_engine):
connection = db_engine.connect()
# begin a non-ORM transaction
transaction = connection.begin()
# bind an individual Session to the connection
db = Session(bind=connection)
# db = Session(db_engine)
yield db
db.rollback()
connection.close()
#
@pytest.fixture(scope="function")
def client(db):
app.dependency_overrides[get_db] = lambda: db
with TestClient(app) as c:
yield cPytest fixtures works like FastAPI dependencies: everything after the yield instruction is ran after exiting the scope pass as a paramter to the decorator. Our db fixture rollsback the session after each test, and we can use it to seed the database. I've also put the dependency override in a fixture alongside the client. That way each
time we use the client, the override will be in place.
Now we can create the fixture we want and test our new endpoint :
@pytest.fixture
def items(db):
create_item(db, schemas.ItemCreate(title="item 1"))
create_item(db, schemas.ItemCreate(title="item 2"))# test_database.py"
def test_list_items(items, client):
response = client.get("/items")
assert len(response.json()) == 2And again our tests pass without issues no matter how many time we run them !
Hooray !
That's it for today, now you know how to properly test you FastAPI app when using a relational database.
If you found this article useful and you want some more follow me on twitter ! 🐦
12