Setting Up a SQLAlchemy and Pytest Based Test Suite
✨ This is an improved article of the previous Python Transactional Tests Using SQL Alchemy, Pytest, and Factory Boy article.
In this post I will show you how you can setup a Pytest test suite that automatically sets up and teardowns your database using SQAlchemy on each test run. It will also provide transactional tests capabilities as described in Python Transactional Tests Using SQL Alchemy, Pytest, and Factory Boy.
Database Setup & Teardown
We are going to be defining the fixtures that will setup and teardown the database. First we will define the database connection responsible for creating the engine and the temporal test database:
# conftest.py
import pytest
from sqlalchemy import create_engine
TEST_DB_NAME = "testdb"
@pytest.fixture(scope="session")
def connection(request):
# Modify this URL according to your database backend
engine = create_engine("mysql+pymysql://root:password@localhost:3306/")
with engine.connect() as connection:
connection.execute(f"CREATE DATABASE {TEST_DB_NAME} CHARACTER SET = 'utf8'")
# Create a new engine/connection that will actually connect
# to the test database we just created. This will be the
# connection used by the test suite run.
engine = create_engine(
f"mysql+pymysql://root:password@localhost:3306/{TEST_DB_NAME}"
)
connection = engine.connect()
def teardown():
connection.execute(f"DROP DATABASE {TEST_DB_NAME}")
connection.close()
request.addfinalizer(teardown)
return connection
This fixture relies on creating two different engines because we first need to create an engine that does not point to any specific database. This is what will allow us to create the actual test database for the test suite execution.
Creating All Database Tables
Now let's go ahead and create the fixture that will create all of our database tables. This assumes that we already have some declared models in our codebase.
# conftest.py
import pytest
from myapp.models import Base
@pytest.fixture(scope="session", autouse=True)
def setup_db(connection, request):
"""Setup test database.
Creates all database tables as declared in SQLAlchemy models,
then proceeds to drop all the created tables after all tests
have finished running.
"""
Base.metadata.bind = connection
Base.metadata.create_all()
def teardown():
Base.metadata.drop_all()
request.addfinalizer(teardown)
This fixture will always run automatically and we do not need to request it thanks to the autouse flag.
Transactional Tests
Even though the test database will be re-created on each Pytest run, we can further improve this and ensure that each test in the suite runs with no data in the database, ensuring a clean state for each test.
# conftest.py
import pytest
# Import a session factory in our app code. Possibly created using
# `sessionmaker()`
from myapp.db import Session
@pytest.fixture(autouse=True)
def session(connection, request):
transaction = connection.begin()
session = Session(bind=connection)
session.begin_nested()
@event.listens_for(session, "after_transaction_end")
def restart_savepoint(db_session, transaction):
if transaction.nested and not transaction._parent.nested:
session.expire_all()
session.begin_nested()
def teardown():
Session.remove()
transaction.rollback()
request.addfinalizer(teardown)
return session
An example of a the session factory in our application could be:
# myapp/db.py
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, scoped_session
engine = create_engine("mysql+pymysql://root:password@localhost:3306/")
Session = scoped_session(sessionmaker(engine))
Notice the use of scoped_session. For more information on this usage, read Contextual/Thread-local Sessions from the SQLAlchemy docs.
Factories
You can easily integrate Factory Boy factories into this workflow:
# tests/factories/base.py
import factory
# Same session factory used by the previous fixture
from myapp.db import Session
class BaseFactory(factory.alchemy.SQLAlchemyModelFactory):
class Meta:
sqlalchemy_session = Session
sqlalchemy_session_persistence = "commit"
Now with an actual factory:
# tests/factories/users.py
from myapp.db.models import User
from tests.factories.base import BaseFactory
class UserFactory(BaseFactory):
class Meta:
abstract = True
model = User
Objects generated by these factories will also be rollbacked in each test thanks to the session fixture we created before. This is
because they both reference the same session factory:
from myapp.db.models import User
from tests.factories.users import UserFactory
class TestUsers:
def test_1(self, session):
user = UserFactory()
assert session.query(User).count() == 1
def test_2(self, session):
assert session.query(User).count() == 0
user = UserFactory.create_batch(3)
assert session.query(User).count() == 3