Python Transactional Tests Using SQL Alchemy, Pytest, and Factory Boy
Knowing how to integrate SQLAlchemy for your automated tests can be a bit tricky. Even after reading several blog posts on how to implement transactional tests with SQLAlchemy and Pytest, it took me a bit of time to make it all play nicely with Factory Boy.
Another key concept that several blog posts out there fail to cover is how SQLAlchemy sessions are handled both in the application and in the test suite.
In this blog post I will teach you how you that is achieved.
What we want to achieve
Basically what we want to achieve with all this is the following:
- Tests live as independent units. They should not care about any existing data in the database, nor about any data that other tests might produce.
- A database rollback should be executed after each test finishes executing. This will ensure a clean state for each test.
- We want to be able to use factories to easily create test objects.
- We do not want any issues and interference between the application database session and the test suite database session.
- We do not want issues in our tests when application code that explicitly calls
session.rollback()
when handling errors is executed.
Database Fixtures
All the magic will happen in some Pytest fixtures that we will define.
Connection
First we will begin with the database connection:
# conftest.py
import pytest
@pytest.fixture(scope="session")
def connection():
db_host = os.getenv("DB_HOST")
db_name = os.getenv("DB_NAME")
db_user = os.getenv("DB_USER")
db_password = os.getenv("DB_PASSWORD")
engine = create_engine(
f"mysql+pymysql://{db_user}:{db_password}@{db_host}:3306/{db_name}"
)
connection = engine.connect()
yield connection
connection.close()
This will create one database connection that will live throught the entire test suite session.
Database Setup
This is an optional step, but comes very handy because you do not need to worry about manually creating tables or running migrations before running your test suite. SQLAlchemy can inspect your SQLAlchemy models to generate the corresponding schema for your tests.
For all this to happen, this fixture needs to be requested explicitly.
# conftest.py
from myproject.models import Base
@pytest.fixture(scope="session")
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)
return None
Session
This is where the real magic happens. This is where we will begin a transaction that will wrap all the tests.
# conftest.py
from sqlalchemy.orm import scoped_session, sessionmaker
from myproject.db import DatabaseManager
# A factory for creating a session here in conftest.py
# It does not matter that our application is also using a
# similar factory to create a session for application use.
Session = scoped_session(sessionmaker())
@pytest.fixture(autouse=True)
def session(connection, setup_db, monkeypatch, request):
transaction = connection.begin()
session = Session(bind=connection)
session.begin_nested()
@event.listens_for(session, "after_transaction_end")
def restart_savepoint(db_session, transaction):
"""Support tests with rollbacks.
This is required for tests that call some services that issue
rollbacks in try-except blocks.
With this event the Session always runs all operations within
the scope of a SAVEPOINT, which is established at the start of
each transaction, so that tests can also rollback the
“transaction” as well while still remaining in the scope of a
larger “transaction” that’s never committed.
"""
if transaction.nested and not transaction._parent.nested:
# ensure that state is expired the way session.commit() at
# the top level normally does
session.expire_all()
session.begin_nested()
def replace_application_session(*args, **kwargs):
# This callable allows us to replace at runtime the
# application's database session with the new test session,
# which supports transactional tests.
return session
# This will depend on how you are declaring and handling the app's
# database session. You will need to monkeypatch accordingly
monkeypatch.setattr(
DatabaseManager,
"session",
replace_application_session
)
def teardown():
Session.remove()
transaction.rollback()
request.addfinalizer(teardown)
return session
This fixture can now be used to perform database operations in our tests in transactional manner:
def test_blah(session):
session.query(MyModel).all()
Moreover, it is automatically called at the beginning of the test suite thanks to autouse=True
. This so that our factories (more on that soon) can easily reference the session after it has been automatically been setup correctly.
This test session fixture is different from all the ones I found in other blog posts because:
- It replaces the application's session (using Pytest's monkeypatch) while still maintaining application functionality ✔
- It will be able to handle parts in the application code that explicity call
rollback()
✔ . Remember that our tests are now being executed inside a transaction!
Adding Factories
The last part of the puzzle is adding factories. And it is super easy, we only need to provide a reference to the session factory used in conftest.py
:
# factories.py
import factory
from myproject.models import MyModel
from tests.conftest import Session
class MyFactory(factory.alchemy.SQLAlchemyModelFactory):
class Meta:
model = MyModel
sqlalchemy_session = Session
sqlalchemy_session_persistence = "commit"
name = factory.Faker("first_name")
And now they can be used in our tests:
from tests.factories import MyFactory
def test_factories():
# You do not need to explicitly request the `session` fixture!
# This object will be rolled back after this test finishes
obj = MyFactory()
assert obj
Conclusion
Coming from a Django background this took me a lot of time to figure out. But I learned a lot in the process and it was definitely worth it! 🎉
References
- https://kampikd.medium.com/setting-up-transactional-tests-with-pytest-and-sqlalchemy-b2d726347629
- https://vittoriocamisa.dev/blog/agile-database-integration-tests-with-python-sqlalchemy-and-factory-boy/
- https://docs.sqlalchemy.org/en/13/orm/session_transaction.html#joining-a-session-into-an-external-transaction-such-as-for-test-suites