Python Transactional Tests Using SQL Alchemy, Pytest, and Factory Boy

Fri, Feb 19, 2021

🆕 Read the newer and improved post Setting Up a SQLAlchemy and Pytest Based Test Suite

📅 Last Updated on November 26 2021

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 1.3 and Pytest, it took me a bit of time to make it all play nicely with Factory Boy.

In this 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 Session Configuration at the Application Level

Your application must provide a place where the database connection is configured so that a session can be used throught the entire application flow. For simplicity we will assume that this lives in a simple module in our application:

# myapp/db.py
import os

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, scoped_session


Session = scoped_session(sessionmaker())

engine = create_engine(
    "mysql+pymysql://{user}:{passwd}@{host}:3306/{db}".format(
        user=os.getenv("DB_USER"),
        passwd=os.getenv("DB_PASSWORD"),
        host=os.getenv("DB_HOST"),
        db=os.getenv("DB_NAME"),
    )
)

Session.configure(bind=get_engine())

Your session configuration will need to be adapted if you require a more dynamic setup, like connecting to different environments using a configuration loader, for example.

✔ Using scoped_session allows disparate sections of the application to call upon a global session, so that all those areas (including tests) may share the same session without the need to pass it explicitly. Because of this, we can simply import this configured Session factory in conftest.py when defining the fixture and in our Factory Boy factories. Read more on scoped_session here.

Database Fixtures

With our application database session configured we can now focus on the Pytest fixtures that will help us achieve transactional tests.

Database Connection

First we will begin with the database connection:

# conftest.py
import pytest

from myapp.db import Session, get_engine


@pytest.fixture(scope="session")
def connection():
    engine = get_engine()
    connection = engine.connect()

    yield connection

    connection.close()

This will create one database connection that will live throught the entire test suite session.

Database Setup (Optional)

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


@pytest.fixture(autouse=True)
def session(connection, setup_db, request):
    """Returns a database session to be used in a test.

    This fixture also alters the application's database
    connection to run in a transactional fashion. This means
    that all tests will run within a transaction, all database
    operations will be rolled back at the end of each test,
    and no test data will be persisted after each test.

    `autouse=True` is used so that session is properly
    initialized at the beginning of the test suite and
    factories can use it automatically.
    """
    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 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 will be able to handle parts in the application code that explicity call rollback() ✔ . Remember that our tests are now being executed inside an explicit 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 myapp.models import MyModel
from myapp.db 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 myapp.models import MyModel
from tests.factories import MyFactory


def test_factories(session):
    # You do not need to explicitly request the `session` fixture!
    # This object will be rolled back after this test finishes
    obj = MyFactory(name="foobar")
    assert obj

    # We can query the database directly for this object as well
    assert session.query(MyModel).filter_by(name="foobar").one_or_none()

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

python sqlalchemy pytest factoryboy