Setting up tests with PostgreSQL and Flask
This is a post describing our test setup at Uptime, using Flask, PostgreSQL and pytest.
The initial problem
Initially, tests within our flask projects were run using a SQLite database, which allowed recreating it between test run, as well as running an in-memory datastore to speed up tests, giving us pretty good performance in our CI pipeline. However, such a solution isn’t viable in the long run: we use PostgreSQL in production, and the two database can behave differently and have different features (in particular, PostgreSQL support the JSON datatype). This then lead us to use a temporary PostgreSQL database for testing, but this comes with more performance issues. The more tests we ran, the slower our CI pipeline got.
The solution described here is our attempt at getting an efficient test setup, that would give us:
- test isolation, to avoid fixtures and db modifications to leak from one test to the other
- good performance
To give a bit of context, our projets currently use:
- python 3.6
- Flask 0.12
- Flask-SQLAlchemy 2.2 / SQLAlchemy 1.1.6
As for the test part, we use Pytest for the runner and Factory-boy to setup test fixtures (I’m not a big fan of using JSON fixtures, for several reasons). Most of our tests use unittest-style syntax for legacy reasons, but the test setup is mostly the same (if only a bit better) using pytest-style syntax. I’ve included a reference project that includes the basis of our flask setup.
Here’s the base class we use for our test suites:
import unittest from project import create_app, database, config def clean_db(db): for table in reversed(db.metadata.sorted_tables): db.session.execute(table.delete()) class BaseTestCase(unittest.TestCase): db = None @classmethod def setUpClass(cls): super(BaseTestCase, cls).setUpClass() cls.app = create_app(app_config=config.TestConfig) cls.db = database.db cls.db.app = cls.app cls.db.create_all() @classmethod def tearDownClass(cls): cls.db.drop_all() super(BaseTestCase, cls).tearDownClass() def setUp(self): super(BaseTestCase, self).setUp() self.client = self.app.test_client() self.app_context = self.app.app_context() self.app_context.push() clean_db(self.db) def tearDown(self): self.db.session.rollback() self.app_context.pop() super(BaseTestCase, self).tearDown()
It works by recreating the database only between test classes (although it could be optimized by doing so only at the start and the end of the test session, using pytest session fixtures for example). In-between tests, the
clean_db function is called to remove any committed object from the database and keep tests isolated.
For reference, here’s the initial naive implementation, that recreated the db between each test:
class BaseTestCase(unittest.TestCase): def setUp(self): super(BaseTestCase, self).setUp() self.app = create_app(app_config=config.TestConfig) self.db = database.db self.db.app = self.app self.db.create_all() self.client = self.app.test_client() self.app_context = self.app.app_context() self.app_context.push() def tearDown(self): self.db.session.rollback() self.db.drop_all() super(BaseTestCase, self).tearDown()
In the end, what kind of speed up did that solution give over the naive implementation? The results aren’t really visible for the dummy project, but here’s what it give for our current test suite:
Naive: 368 passed, 3 skipped, 10 warnings in 360.57 seconds Improved: 368 passed, 3 skipped, 10 warnings in 182.19 seconds
So overall a 50% speedup!
If you’d like to see more about test setup, or a bigger flask project in general, have a look at the skylines project, which I found very well architected and which gave me some reference to rethink how our tests should be setup.