Flask-SQLAlchemy and PostgreSQL Unit Testing with Transaction Savepoints

5.00 avg. rating (95% score) - 1 vote

PostgreSQL provides us two very powerful features which are helpful with unit testing: transactional DDL and transaction savepoints. In this article I will show you how to use those with Flask-SQLAlchemy unit tests.

Transactional DDL means you can create tables inside a transaction, run tests against them, and roll back your changes after you are done. The database will be left in the same state as it was when you started. If you started with an empty database, nothing will be left in the database afterwards.

Savepoints allow you to roll back a part of a transaction without affecting what has happened before that savepoint was created. You can run a single test inside a savepoint and roll back just the changes that single test made without affecting the changes your set-up code has done before the savepoint.

That means you can create a large number of tables and other database objects in the beginning of your test suite and then run individual tests inside nested transaction using savepoints. There is no need to drop and re-create the whole database schema for each test.

PostgreSQL Database for Unit Tests

We will need an empty database to run our unit tests against. If you have PostgreSQL already installed, at least on Ubuntu there are command line tools to get it done:

The createuser command creates a database user. The createdb creates the database, and the -O option sets the owner of the database.

Other Requirements

You obviously need to install Flask-SQLAlchemy and psycopg2 for PostgreSQL connectivity:

An App, a Model, and Some Tests

First let’s take a look at an ordinary test suite. Flask-SQLAlchemy wraps SQLAlchemy and integrates it with Flask, providing some additional functionality with it. A very minimal Flask app with Flask-SQLAlchemy and PostgreSQL support could look something like below.

Please note that I haven’t tested the following example nor have I run its tests. The downloadable file at the end of this article contains a working and tested example, and includes the tests used to test it. The example below gives us something we can start do discuss here. Let me know if you find errors.

Our app is in deed very minimal. It doesn’t even have any views. But it is enough for our purposes.

A minimal test suite for our database model might look like this:

Note that the db.create_all/drop_all is executed for each test, as well as the initial model adding. We can avoid this by using a transaction and some savepoints.

Tests in Nested Transactions

So how do we use savepoints to our advantage when writing unit tests for Flask-SQLAlchemy apps? To test arbitrary Flask-SQLAlchemy code inside a nested transaction, we must replace the global db.session object with one that runs in a nested transaction. After all, your code is likely to perform all of its database operations through the db.session object.

We can choose to use either the TestCase class’s setUpClass/tearDownClass classmethods which are called once per each TestCase, or the setUpModule/tearDownModule functions which unittest will call once per module when running the tests. And we can even use both if we want to do some setup work at the module level and some more in each class.

Here’s an example of how to do it using the classmethods:

The setUpClass method is only run once in the beginning of the test. The database schema is created there. The schema is dropped in the end in the tearDownClass method.

The setUp method now creates two savepoints. Strictly speaking only one savepoint is needed, but I would recommend two because the code you are testing might call db.session.commit() and that will commit the savepoint it is running in. A committed savepoint can not be rolled back to. So in the tearDown method we will roll back to the first savepoint which will be available also in the case the tested code commits. When the savepoint is rolled back, the next test case will see the database exactly like it was after the setUpClass call.

Note that if your code does multiple commits per test, this doesn’t work and you will receive exceptions from SQLAlchemy. You can add more savepoints to work around that. You must have one more savepoint than you have commits in a single test. On the other hand, it is a good way of finding code paths that perform multiple commits. Doing lots of commits will result in a lot of disk activity and will slow down your application.

If you want to do multiple commits in your test case, you can use db.session.begin_nested and commit that:

Stuff you do inside a committed savepoint will be available later on, so the savepoint1 stuff in the example above will be available inside savepoint2, unless you choose to roll back to savepoint1 instead of committing. Eventually of course the test tearDown will roll back to its own savepoint to undo everything you did before running the next test.

Transactional DDL

The next logical step is to run db.create_all() inside the same transaction as the tests. This way the database you run unit tests against will start empty and end up empty without even the need to drop tables afterwards.

In fact it might even be possible to run it against a database which has existing tables already if you drop first, and keep your old data after the tests. But I haven’t tested if it works or not. It is probably not a good idea, though. There might be side effects since some database objects, such as sequences, are not kept completely isolated inside the transaction.

Flask-SQLAlchemy does not allow us to run the metadata creation inside the same transaction as the rest of the app. Again we need to monkey patch the db object and replace the session with our own, this time to make a new database connection. Flask-SQAlchemy depends on the db.session object to be a factory for new sessions, but we want to avoid new sessions since we want everything to run inside a single transaction. We will have to use a custom session which is a “factory” that returns itself.

Below is an example of how to do it in the setUpModule/tearDownModule functions.

What setUpModule does is it makes its own database connection instead of using the one Flask-SQLAlchemy has created. This is needed so that we can begin a transaction, run the DDL in it, and continue with running the tests in the same transaction. The Flask-SQLAlchemy db.create_all() will commit the schema into the database if we use it directly. The TestingSession is a modified Flask-SQLAlchemy SignallingSessionclass. It keeps everything running in our single transaction.

Another detail to note is the db.session.remove = lambda: None statement in the test setUp code. It is a required fix when you push the application context in your own setUp method and pop it in the tearDown method. Flask-SQLAlchemy will raise an exception during app context teardown without it.

The rest of it is a simple extension of the original idea. Here we have two levels of setup code instead of one, and savepoints to guard them both. The setUpModule call creates the schema once for the whole module, while the model creation is done in the setUpClass method. It means that if you add another TestCase in the same module the schema will stay intact but the model inserted or any other operation performed in MyModelTestCase’s setUpClass will only be there for MyModelTestCase’s tests.

Abstract it Out

It is obviously not a good idea to keep code like that in each test module so let’s extract the functionality and put it in a separate module which we will call flask_sqlalchemy_testing. The code for this and the tests can be downloaded at the end of the article.

Let’s use an app which has two models for a change:

Here’s the test module:

You can run the tests with “python test_app.py”. You can uncomment the lines which enable sqlalchemy.engine logging to see the SQL it is issuing.

We must of course call the setUpModule of the flask_sqlalchemy_testing module in the start of our own setUpModule. If you don’t need to override it, remember to import it. The tearDownModule also needs to be imported so that it will run. The same goes for setUpClass/tearDownClass.

Download the Code

The code is available to download and use as you wish. Here it is:

flask_sa_testing.tar.gz

I have used it, but not very extensively yet. Let me know if you find errors in it.

Here are the commands to run it in a virtualenv:

 

5 thoughts on “Flask-SQLAlchemy and PostgreSQL Unit Testing with Transaction Savepoints”

  1. Here’s an additional context manager you can use to capture commits in your test code so that they will not be committed to the unittest database. With this it is possible to handle multiple commits.

    It will make the code commit a savepoint instead of the full transaction. The transaction can be later rolled back, rolling back the savepoint as well. You can also give it a rollback=True argument which will roll back immediately.

    Use it like this:

  2. Thank you for your work on this subject. I think it’s currrenty a real issue
    for performance of DB tests.

    I tried your solution but I found two issues with external modules:
    – FactoryBoy doesn’t use the right session:
    https://github.com/rbarrois/factory_boy/blob/master/factory/alchemy.py#L44

    I didn’t know the proper fix so I overriden the class to remove the session.add()
    and do it explicitly in my test.

    – sqlalchemy-utils provides ChoiceType for example, if you define something like
    class Menu(db.Model):
    FRUIT_CHOICES = [
    (‘A’, “Apple”),
    (‘B’, “Banana”),
    ]
    fruit = db.Column(ChoiceType(FRUIT_CHOICES))

    you set fruit to ‘A’, it’s only when the Menu instance is commited the fruit field if updated
    to provide menu.fruit.value -> “Apple” but this behaviour seems to don’t work with the savepoint, so
    for now, I used a session.flush() and session.expire(menu) to workaround it…

    I would be great if you are willing to talk about your project on the Flask or SQLAlchmey mailing lists.
    The subject has already been discussed https://groups.google.com/forum/#!msg/sqlalchemy/4cJWKInhNm8/O4jTGx0hUUwJ

    1. I just do something like this:

      My create_app expects a config name as argument which in this case is ‘unittest’.

Leave a Reply