Rollback transactions not working with py.test and Flask

Question

I'm using py.test to test my Flask application but I am getting IntegrityError because I am creating the same model in two different tests.

I am using postgreSQL 9.3.5 and Flask-SQLAlchemy 1.0.

EDIT I've updated my sessoin fixture with Jeremy Allen's answer and it fixed a lot of the errors. However it seems when I use the user fixture that I still get IntegrityErrors

Error

E       IntegrityError: (IntegrityError) duplicate key value violates unique constraint "ix_users_email"
E       DETAIL:  Key (email)=([email protected]) already exists.
E        'INSERT INTO users (email, username, name, role_id, company_id, password_hash, confirmed, member_since, last_seen) VALUES (%(email)s, %(username)s, %(name)s, %(role_id)s, %(company_id)s, %(password_hash)s, %(confirmed)s, %(member_since)s, %(last_seen)s) RETURNING users.id' {'username': 'not_used', 'confirmed': True, 'name': 'To be Removed', 'member_since': datetime.datetime(2014, 10, 29, 19, 19, 41, 7929), 'company_id': None, 'role_id': 3, 'last_seen': datetime.datetime(2014, 10, 29, 19, 19, 41, 7941), 'email': '[email protected]', 'password_hash': 'pbkdf2:sha1:1000$cXUh6GbJ$6f38242871cff5e4cce4c1dc49a62c4aea4ba1f3'}

conftest.py

@pytest.yield_fixture(scope='session')
def app():
    app = create_app('testing')
    app.config['SERVER_NAME'] = 'example.com:1234'
    ctx = app.app_context()
    ctx.push()
    app.response_class = TestResponse
    app.test_client_class = TestClient
    yield app
    ctx.pop()


@pytest.fixture(scope='session')
def db(app):
    _db.drop_all()
    _db.create_all()

    Permission.insert_initial()
    Role.insert_initial()
    Technology.insert_initial()
    Product.insert_initial()
    Actor.insert_initial()
    Industry.insert_initial()
    DeliveryCategory.insert_initial()
    DeliveryMethod.insert_initial()

    user = User(email='[email protected]', username='admin', confirmed=True, password='admin', name='Admin')
    user.role = Role.query.filter_by(name='Administrator').first()
    _db.session.add(user)
    _db.session.commit()

    return _db


@pytest.yield_fixture(scope='function')
def session(db):
    db.session.begin_nested()
    yield db.session
    db.session.rollback()


@pytest.yield_fixture(scope='function')
def user(session):
    yield session.query(User).filter_by(email='[email protected]').first()


@pytest.yield_fixture(scope='function')
def client(app, user):
    client = app.test_client()
    client.auth = 'Basic ' + b64encode((user.email + ':' + 'admin').encode('utf-8')).decode('utf-8')
    yield client

Tests that fail

def test_edit_agenda_add_company_rep_without_company(session, client, user):
    user2 = User(name='To be Removed', password='not_used', username='not_used', confirmed=True,
                email='[email protected]', role=Role.query.filter_by(name='User').first())
    agenda = Agenda(name='Invalid Company Rep', creator=user)
    session.add(agenda)
    session.commit()

    response = client.jput('/api/v1.0/agendas/%s' % agenda.id,
        data={
            'company_representative': user2.id
        }
    )
    assert response.status_code == 200

def test_edit_agenda_add_user_already_in_agenda(session, client, user):
    user2 = User(name='To be Removed', password='not_used', username='not_used', confirmed=True,
                email='[email protected]', role=Role.query.filter_by(name='User').first())
    agenda = Agenda(name='Invalid Company Rep', creator=user)
    agenda.users.append(user2)
    session.add(agenda)
    session.commit()

    response = client.jput('/api/v1.0/agendas/%s' % agenda.id,
        data={
            'users': [user2.id]
        }
    )
    assert response.status_code == 200

Tests that pass

def test_get_agenda_modules_where_agenda_that_does_not_exist(session, app):
    # Create admin user with permission to create models
    user = User(email='[email protected]', username='admin2', confirmed=True, password='admin2')
    user.role = Role.query.filter_by(name='Administrator').first()
    session.add(user)
    session.commit()

    client = app.test_client()
    client.auth = 'Basic ' + b64encode(
        (user.email + ':' + 'admin2').encode('utf-8')).decode('utf-8')
    response = client.jget('/api/v1.0/agenda-modules/%s/%s' % (5, 4))
    assert response.status_code == 404

def test_get_agenda_modules_agenda_modules_does_not_exist(session, app):
    agenda = Agenda(name='Is tired in the AM')
    session.add(agenda)

    # Create admin user with permission to create models
    user = User(email='[email protected]', username='admin2', confirmed=True, password='admin2')
    user.role = Role.query.filter_by(name='Administrator').first()
    session.add(user)
    session.commit()

    client = app.test_client()
    client.auth = 'Basic ' + b64encode(
        (user.email + ':' + 'admin2').encode('utf-8')).decode('utf-8')
    response = client.jget('/api/v1.0/agenda-modules/%s/%s' % (agenda.id, 4))
    assert response.status_code == 400
    assert response.jdata['message'] == 'AgendaModule does not exist.'

Answer 1

It looks like you are trying to join a Session into an External Transaction and you are using flask-sqlalchemy.

Your code is not working as expected because the session actually ends up using a different connection to the one you are beginning the transaction on.

1. You need to bind the Session to the Connection

As in the example linked above. A quick change to your code in conftest.py should do it:

@pytest.yield_fixture(scope='function')
def session(db):
    ...
    session = db.create_scoped_session(options={'bind':connection})
    ...

Unfortunately, due to flask-sqlalchemy's SignallingSession (as at v2.0), your 'bind' argument will be overruled!

This is because SignallingSession set the 'binds' argument such that it will take precedence over our 'bind' argument and it doesn't offer us a nice way to specify our own 'binds' argument.

There is a GitHub pull request from December 2013 where someone else had the same problem.

2. Tweak flask-sqlalchemy

We can subclass SignallingSession to allow us to do what we want:

class SessionWithBinds(SignallingSession):
    """The extends the flask-sqlalchemy signalling session so that we may
    provide our own 'binds' argument.
    """

    def __init__(self, db, autocommit=False, autoflush=True, **options):
        #: The application that this session belongs to.
        self.app = db.get_app()
        self._model_changes = {}
        #: A flag that controls whether this session should keep track of
        #: model modifications.  The default value for this attribute
        #: is set from the ``SQLALCHEMY_TRACK_MODIFICATIONS`` config
        #: key.
        self.emit_modification_signals = \
            self.app.config['SQLALCHEMY_TRACK_MODIFICATIONS']
        bind = options.pop('bind', None) or db.engine
        # Our changes to allow a 'binds' argument
        try:
            binds = options.pop('binds')
        except KeyError:
            binds = db.get_binds(self.app)
        SessionBase.__init__(self, autocommit=autocommit, autoflush=autoflush,
                             bind=bind,
                             binds=binds, **options)

And then subclass SQLAlchemy (the main flask-sqlalchemy class) to use our SessionWithBinds in place of SignallingSession

class TestFriendlySQLAlchemy(SQLAlchemy):
    """For overriding create_session to return our own Session class"""

    def create_session(self, options):
        return SessionWithBinds(self, **options)

Now you have to use this class in place of SQLAlchemy:

db = TestFriendlySQLAlchemy()

And finally, back in our conftest.py specify a new 'binds':

@pytest.yield_fixture(scope='function')
def session(db):
    ...
    session = db.create_scoped_session(options={'bind':connection, 'binds':None})
    ...

Now your transactions should rollback as expected.

This is all a bit complicated...

Instead of doing all this you could try using Session.begin_nested. It requires that your database supports SQL SAVEPOINTs (PostgreSQL does).

Change your conftest.py fixture:

@pytest.yield_fixture(scope='function')
def session(db):
    db.session.begin_nested()
    yield db.session
    db.session.rollback()

More info on using SAVEPOINTs in SQLAlchemy: http://docs.sqlalchemy.org/en/latest/orm/session_transaction.html#using-savepoint

This is pretty simple but will work as long as the code you are testing doesn't call rollback itself. If this is a problem take a look at the code here in the SQLAlchemy docs under the title "Supporting Tests with Rollbacks"

Answer 2

The key here is to run your tests within a nested session, and then rollback everything after the execution of each test (this also assumes there are no dependencies across your tests).

I suggest adopting the following approach, by running each of your tests within a nested transaction:

# module conftest.py
import pytest

from app import create_app
from app import db as _db
from sqlalchemy import event
from sqlalchemy.orm import sessionmaker

@pytest.fixture(scope="session")
def app(request):
    """
    Returns session-wide application.
    """
    return create_app("testing")


@pytest.fixture(scope="session")
def db(app, request):
    """
    Returns session-wide initialised database.
    """
    with app.app_context():
        _db.drop_all()
        _db.create_all()


@pytest.fixture(scope="function", autouse=True)
def session(app, db, request):
    """
    Returns function-scoped session.
    """
    with app.app_context():
        conn = _db.engine.connect()
        txn = conn.begin()

        options = dict(bind=conn, binds={})
        sess = _db.create_scoped_session(options=options)

        # establish  a SAVEPOINT just before beginning the test
        # (http://docs.sqlalchemy.org/en/latest/orm/session_transaction.html#using-savepoint)
        sess.begin_nested()

        @event.listens_for(sess(), 'after_transaction_end')
        def restart_savepoint(sess2, trans):
            # Detecting whether this is indeed the nested transaction of the test
            if trans.nested and not trans._parent.nested:
                # The test should have normally called session.commit(),
                # but to be safe we explicitly expire the session
                sess2.expire_all()
                sess2.begin_nested()

        _db.session = sess
        yield sess

        # Cleanup
        sess.remove()
        # This instruction rollsback any commit that were executed in the tests.
        txn.rollback()
        conn.close()