Disclaimer:
I’m a total newb with Python and the ecosystem that surrounds it. I come from Rails and apparently I’ve been completely spoiled by how easy it was to use.
Dear reader, I’m writing to declare victory. I’ve successfully navigated the gauntlet of setting up a test database for my app that consists of FastAPI, SQLAlchemy, Alembic, and PyTest with PostgreSQL as the database. In my whole career, I’ve never actually had to do this seemingly boilerplate task. Since documentation was scarce, I’m leaving this breadcrumb in the hopes that it will benefit those that come after me.
Step By Step
Step 1: Forget about Alembic. Yes, you used it to get your database schema to this point, but you won’t need it for this exercise. Save yourself the time. Don’t bark up that tree.
Step 2: Drop and re-create the test database from previous runs. This goes in conftest.py , which gets loaded at the beginning of a test run. You’ll see some references to settings.db_url . We’ll cover that part in a moment. The main idea here is that 1) we set up DB_NAME to be something other than our normal development database, and 2) we call db_prep, which drops (if it existed previously) and re-creates the test database and adds our apps user to have access to the test database. It’s just executing the same SQL that you would execute if you were doing this by hand. We’ll cover the fixture at the bottom of this file that invokes prep_db in just a minute.
# test
# |_ conftest.pyimport os
from dotenv import load_dotenv
from sqlalchemy import create_engine, MetaData
from sqlalchemy.exc import ProgrammingError, OperationalError
import pytest
from app import settingsload_dotenv()
settings.init()DB_NAME = f”{os.getenv(‘DB_NAME’)}_test”
settings.db_url = f”postgresql://{os.getenv(‘DB_USER’)}:{os.getenv(‘DB_PASSWORD’)}@{os.getenv(‘DB_HOST’, ‘localhost’)}/{DB_NAME}”def db_prep():
print(“dropping the old test db…”)
engine = create_engine(“postgresql://postgres@localhost/postgres”)
conn = engine.connect()
try:
conn = conn.execution_options(autocommit=False)
conn.execute(“ROLLBACK”)
conn.execute(f”DROP DATABASE {DB_NAME}”)
except ProgrammingError:
print(“Could not drop the database, probably does not exist.”)
conn.execute(“ROLLBACK”)
except OperationalError:
print(“Could not drop database because it’s being accessed by other users (psql prompt open?)”)
conn.execute(“ROLLBACK”)
print(f”test db dropped! about to create {DB_NAME}”)
conn.execute(f”CREATE DATABASE {DB_NAME}”)
try:
conn.execute(f”create user {os.getenv(‘DB_USER’)} with encrypted password ‘{os.getenv(‘DB_PASSWORD’)}’”)
except:
print(“User already exists.”)
conn.execute(f”grant all privileges on database {DB_NAME} to {os.getenv(‘DB_USER’)}”)
conn.close()
print(“test db created”)@pytest.fixture(scope=”session”, autouse=True)
def fake_db():
db_prep()
print(f”initializing {DB_NAME}…”)
engine = create_engine(settings.db_url)
from app.models import User
from app.database import SessionLocal, Base
db = SessionLocal()
Base.metadata.create_all(engine)
print(f”{DB_NAME} ready to rock!”)
try:
yield db
finally:
db.close()
Step 3: I created a settings file that helps me manage global variables. In this case, we’re using it to make our connection string to the database global in scope. You’ll see why in step 4.
# app
# |_ settings.pydef init():
global db_url
db_url = None
Step 4: Use the database connection string from settings if it exists otherwise initialize it. This allows us to initialize the database name to our test database if we’re entering the app through tests. Otherwise we initialize it in main to our dev/prod database name. NOTE: We want to make sure we import anything dependent on the database after we set the global db_url setting.
# app
# |_ main.pyimport os
import traceback
from typing import Optional
from dotenv import load_dotenv
from fastapi import Depends, FastAPI, Header, HTTPException, Request
from fastapi.responses import PlainTextResponse
from sqlalchemy import text
from sqlalchemy.orm import Session
import requests
from . import settingsload_dotenv()try:
settings.db_url
except AttributeError:
settings.init()
settings.db_url = f"postgresql://{os.getenv('DB_USER')}:{os.getenv('DB_PASSWORD')}@{os.getenv('DB_HOST', 'localhost')}/{os.getenv('DB_NAME')}"from .database import SessionLocal, engine
from . import models, schemasapp = FastAPI()def get_db():
db = SessionLocal()
try:
yield db
finally:
db.close()
And in our database set-up, we only reference the setting.
# app
# |_ database.pyimport os
from dotenv import load_dotenv
from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from . import settingsengine = create_engine(settings.db_url)SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)Base = declarative_base()
Step 5: OK, let’s take a look back at the fixture from conftest.py :
@pytest.fixture(scope=”session”, autouse=True)
def fake_db():
db_prep()
print(f”initializing {DB_NAME}…”)
engine = create_engine(settings.db_url)
from app.models import User
from app.database import SessionLocal, Base
db = SessionLocal()
Base.metadata.create_all(engine)
print(f”{DB_NAME} ready to rock!”)
try:
yield db
finally:
db.close()Here’s what’s going on, line-by-line:
1. The fixture is declared with session scope.
2. We call the afore-referenced db_prep to drop the test test database if it existed before and re-create it.
3. We create our SQLAlchemy engine using settings.db_url, which was set up to point at our test database.
4. We import any models that define our schema. In this case it’s just the User model. It’s important that we do this at this point. (I’m assuming at this point that you’ve created SQLAlchemy models and such.)
5. We import other SQLAlchemy dependencies for database access during tests.
6. We call Base.metatdata.create_all and pass it the engine we instantiated earlier that points to our test database. This does the work of inspecting any models that have been imported and creating the schema that represents them in our test database.
7. We yield the database session we’ve created to allow our tests to use it. (See below.)
Step 6: Use the test database in your tests:
import pytest
from app.models import User
from app.schemas import UserCreateclass TestModels:
def test_user(self, fake_db):
user = UserCreate(external_id="890", external_id_source="456")
u = User.create(db=fake_db, user=user)
assert u.id > 0
The fake_db argument in our test is a reference to the fixture we created above. My model has a create method that takes the database connection as a dependency to be injected. This may be refactored shortly, but it works for the moment.
That’s All Folks
I’m 100% positive that somebody else has a better way of doing this or will see a better way to do this based on what I’ve shared. At least I hope that’s the case. If you’re that somebody and happen to come across this, give me a holler and show me (and everyone else) the way. Like I said at the outset, I’m no expert with Python yet. I learned a lot in this process and I’m always open to learning from those that have gone before me.
