Running alembic migrations in flask tests

I was recently writing some tests for one of my side projects. It is a traditional API that is built on top of flask and alembic (or more specifically flask-migrate).

One of the things that makes it a little more special is that it makes use of custom PostgreSQL functions at parts. Those functions are also managed by alembic. I essentially have some migrations in which I just execute CREATE OR REPLACE FUNCTION ...-SQL statements.

When I was writing tests I wanted to test those functions as well so I needed a way to get them into the test database. The usual and widespread way of getting your database to up to date with your models is to use db.create_all(). This was not possible for my use case though as it only creates what is inside the models but does not find the custom functions as those only live in alembic.

So I needed a way to execute all the migrations. It is actually not very difficult and all parts to this solution can be found somewhere on the web but I did not find it written down as a complete example so I thought I would share my version:

import unittest
import os

from flask_migrate import upgrade

from myapi.app import app, db

class DBTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        app.config['SQLALCHEMY_DATABASE_URI'] = os.environ.get("DATABASE_URL_TEST")
        app.config['TESTING'] = True
        with app.app_context():
            upgrade()

    @classmethod
    def tearDownClass(cls):
        db.drop_all()
        db.engine.execute("DROP TABLE alembic_version")

    def setUp(self):
        self.app = app.test_client()

    def tearDown(self):
        for table in reversed(db.metadata.sorted_tables):
            db.engine.execute(table.delete())
        db.session.commit()
        db.session.remove()

Here is what is going on: The setUpClass method is called once per test class. It makes sure to set up the test database (so that we do not accidentally destroy our development/production database) and sets flask to test mode. It also does the interesting part of executing the migrations. To do that it needs a so called flask app_context but as you can see this is pretty easy to get with a with statement.

The tearDownClass is executed after all the test of this class have been executed. It uses sqlalchemy to get all the tables that it knows about and drops those so that we have an empty database for the next test again. We also need to manually drop the alembic_version table. This is where alembic remembers at which revision the database model should be. If we do not reset this the next time that we run the tests it would check the table, see that there is an entry pointing to the latest revision and not do anything. So we drop that table as well to make sure that we start afresh.

The setUp class which is executed before each test actually does nothing particularly interesting regarding the database. It just provides an app that can be used in the tests later.

Finally the tearDown function which is executed after every test deletes all entries from all databases so that the test can be run in any order and do not interfere with each other. This uses sqlalchemy’s cool MetaData.sorted_tables function which gives you a list of all your models/tables taking into account their relationships. We reverse this list so that we get the tables without dependencies first and drop them one after the other. Finally we close the session and are good to go for the next step.

We can now use this as the base class / mixin for our tests by extending it in our test classes like this:

from helpers import DBTest

class TestController(DBTest):
  def test_addition(self):
    pass
    # ...

This will make sure that for every test class that we write we get a freshly migrated database and for every test in there the database is going to be empty. Since this includes all alembic migrations this is also going to give us the custom functions or anything else that we defined in alembic that is not mirrored in our models.

I hope that this helps you if you find yourself in a similar situation. If there is something that does not work, room for improvement or questions feel free to contact me by twitter or email.