Creating a Test Database: PyTest + SQLAlchemy

Step By Step


# test
# |_ conftest.py
import os
from dotenv import load_dotenv
from sqlalchemy import create_engine, MetaData
from sqlalchemy.exc import ProgrammingError, OperationalError
import pytest
from app import settings
load_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()
# app
# |_ settings.py
def init():
global db_url
db_url = None
# app
# |_ main.py
import 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 settings
load_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, schemas
app = FastAPI()def get_db():
db = SessionLocal()
try:
yield db
finally:
db.close()
# app
# |_ database.py
import 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 settings
engine = create_engine(settings.db_url)SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)Base = declarative_base()
@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()
import pytest
from app.models import User
from app.schemas import UserCreate
class 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

That’s All Folks

--

--

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store