How to combine py.test fixtures with Flask-SQLAlchemy and PostgreSQL?

Question

I'm struggling to write py.test fixtures for managing my app's database that maximize speed, supports pytest-xdist parallelization of tests, and isolates the tests from each other.

I'm using Flask-SQLAlchemy 2.1 against a PostgreSQL 9.4 database.

Here's the general outline of what I'm trying to accomplish:

1. $ py.test -n 3 spins up three test sessions for running tests.

2. Within each session, a py.test fixture runs once to setup a transaction, create the database tables, and then at the end of the session it rolls back the transaction. Creating the database tables needs to happen within a PostgreSQL transaction that's only visible to that particular test-session, otherwise the parallelized test sessions created by pytest-xdist cause conflicts with each other.

3. A second py.test fixture that runs for every test connects to the existing transaction in order to see the created tables, creates a nested savepoint, runs the test, then rolls back to the nested savepoint.

4. Ideally, these pytest fixtures support tests that call db.session.rollback(). There's a potential recipe for accomplishing this at the bottom of this SQLAlchemy doc.

5. Ideally the pytest fixtures should yield the db object, not just the session so that folks can write tests without having to remember to use a session that's different than the standard db.session they use throughout the app.

Here's what I have so far:

import pytest

# create_app() is my Flask application factory
# db is just 'db = SQLAlchemy()' + 'db.init_app(app)' within the create_app() function
from app import create_app, db as _db 


@pytest.yield_fixture(scope='session', autouse=True)
def app():
    '''Session-wide test application'''
    a = create_app('testing')
    with a.app_context():
        yield a

@pytest.yield_fixture(scope='session')
def db_tables(app):
    '''Session-wide test database'''
    connection = _db.engine.connect()
    trans = connection.begin() # begin a non-ORM transaction

    # Theoretically this creates the tables within the transaction
    _db.create_all()
    yield _db
    trans.rollback()
    connection.close()

@pytest.yield_fixture(scope='function')
def db(db_tables):
    '''db session that is joined to existing transaction'''

    # I am quite sure this is broken, but it's the general idea 

    # bind an individual Session to the existing transaction
    db_tables.session = db_tables.Session(bind=db_tables.connection)

    # start the session in a SAVEPOINT...
    db_tables.session.begin_nested()

    # yield the db object, not just the session so that tests
    # can be written transparently using the db object
    # without requiring someone to understand the intricacies of these
    # py.test fixtures or having to remember when to use a session that's
    # different than db.session
    yield db_tables

    # rollback to the savepoint before the test ran
    db_tables.session.rollback()
    db_tables.session.remove() # not sure this is needed

Here's the most useful references that I've found while googling:

http://docs.sqlalchemy.org/en/latest/orm/session_transaction.html#joining-a-session-into-an-external-transaction-such-as-for-test-suites

http://koo.fi/blog/2015/10/22/flask-sqlalchemy-and-postgresql-unit-testing-with-transaction-savepoints/

https://github.com/mitsuhiko/flask-sqlalchemy/pull/249

Answer 1

I'm a couple years late here, but you might be interested in pytest-flask-sqlalchemy, a plugin I wrote to help address this exact problem.

The plugin provides two fixtures, db_session and db_engine, which you can use like regular Session and Engine objects to run updates that will get rolled back at the end of the test. It also exposes a few configuration directives (mocked-engines and mocked-sessions) that will mock out connectables in your app and replace them with these fixtures so that you can run methods and be sure that any state changes will get cleaned up when the test exits.

The plugin should work with a variety of databases, but it's been tested most heavily against Postgres 9.6 and is in production in the test suite for https://dedupe.io. You can find some examples in the documentation that should help you get started, but if you're willing to provide some code I'd be happy to demonstrate how to use the plugin, too.