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

Fri, Feb 19, 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 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

python sqlalchemy pytest factoryboy