Wrap Flask/Sqlalchemy Tests in a Transaction

So, I wanted to share a little bit of information that I have accumulated. I am sure that a lot of other people know how to do this, since everything here is pulled from a variety of places, but I have never seen it all put together. The basic issue is if you are writing a Flask application, using SqlAlchemy, unit testing can be a bit hard. Hopefully you are using Flask-Sqlalchemy and Flask-Testing which are designed to simplify some of this integration. However, the recommended way of running your tests is to either use an in-memory sqlite db (which can lead to problems of code that pasts unit tests but fails with the real database), or to drop and recreate your database in between every test. As you might imagine, the second suggestion is very slow :)

The ideal, for me at least, is to wrap each test case in a transaction so that you can just do rollback and get back to the initial state. However, you have to worry about all the different places where the code under test may call commit! The solution is nested transactions. Not all database systems will work with this method, and not being a DBA I could not even guess which ones I am talking about... it DOES work with Postgresql and Mysql though, so that should satisfy a pretty large percentage of use cases!

Flask Testing helps you create a base TestCase that handles creating your Flask App and contexts and tearing it down in between tests. This builds on that a little bit- any test case that inherits from this Class will start an outer db transaction before your setup method, then map db.session to an inner transaction... and set a hook so that anytime someone commits or rollbacks that transaction, it gets restarted automatically. The transaction is then disposed of after your tearDown method is called. The code!

from flask_testing import TestCase 
# I could recreate the whole app inside the setup, but I have seen this have a drastic 
# improvement in speed, and generally don't put any logic that can be polluted on the app.
# your mileage may vary here
app = create_app() # This inherits from flask_testing
class BaseCase(TestCase): 
    def create_app(self): 
        return app 
    # override pre_setup and post_teardown to add our transaction logic in the correct 
    # location 
    def _pre_setup(self): 
        super(BaseCase, self)._pre_setup() 
        self._start_transaction() 
    def _post_teardown(self): 
        try: 
            self._close_transaction() 
        finally: 
            super(BaseCase, self)._post_teardown() 
    def _start_transaction(self): 
        # Create a db session outside of the ORM that we can roll back 
        self.connection = db.engine.connect() 
        self.trans = self.connection.begin() 
        # bind db.session to that connection, and start a nested transaction 
        db.session = db.create_scoped_session(options={'bind': self.connection}) 
        db.session.begin_nested() 
        # sets a listener on db.session so that whenever the transaction ends- 
        # commit() or rollback() - it restarts the nested transaction 
        @event.listens_for(db.session, "after_transaction_end") 
        def restart_savepoint(session, transaction): 
            if transaction.nested and not transaction._parent.nested: 
                session.begin_nested() 
                self._after_transaction_end_listener = restart_savepoint 
        def _close_transaction(self): 
            # Remove listener 
            event.remove(db.session, "after_transaction_end", self._after_transaction_end_listener) 
            # Roll back the open transaction and return the db connection to 
            # the pool 
            db.session.close() 
            db.get_engine(self.app).dispose() 
            self.trans.rollback() 
            self.connection.invalidate()

An important caveat to keep in mind, this runs right before each test methods setUp function. This means that anything you do in setupClass or the like will not be wrapped in a transaction. Finally, the Flask-Sqlalchemy app currently has a bug that blows up the session event hooks. The development version has this fixed, but hasn't been released as of today. I use

git+https://github.com/mitsuhiko/flask-sqlalchemy.git#egg=Flask-SQLAlchemy

in my requirements.txt to get that version. The issue was fixed in this PR, and should be in the next release after 2.1, whenever that is. You could also fix this by inheriting from the base class and implementing the fix yourself if you prefer.

Good luck!