Writing Unit Tests for Django Migrations

Editor's note: This post was originally published in February 2016 and was updated in August 2017 to incorporate improvements suggested by our readers. It has also been tested for compatibility as of the Django 1.11 release.

Testing in a Django project ensures the latest version of a project is as bug-free as possible. But when deploying, you’re dealing with multiple versions of the project through the migrations.

The test runner is extremely helpful in its creation and cleanup of a test database for our test suite. In this temporary test database, all of the project's migrations are run before our tests. This means our tests are running the latest version of the schema and are unable to verify the behavior of those very migrations because the tests cannot set up data before the migrations run or assert conditions about them.

We can teach our tests to run against those migrations with just a bit of work. This is especially helpful for migrations that are going to include significant alterations to existing data.

The Django test runner begins each run by creating a new database and running all migrations in it. This ensures that every test is running against the current schema the project expects, but we'll need to work around this setup in order to test those migrations. To accomplish this, we'll need to have the test runner step back in the migration chain just for the tests against them.

Ultimately, we're going to try to write tests against migrations that look like this:

class TagsTestCase(TestMigrations):

    migrate_from = '0009_previous_migration'
    migrate_to = '0010_migration_being_tested'

    def setUpBeforeMigration(self, apps):
        BlogPost = apps.get_model('blog', 'Post')
        self.post_id = BlogPost.objects.create(
            title = "A test post with tags",
            body = "",
            tags = "tag1 tag2",
        ).id

    def test_tags_migrated(self):
        BlogPost = self.apps.get_model('blog', 'Post')
        post = BlogPost.objects.get(id=self.post_id)

        self.assertEqual(post.tags.count(), 2)
        self.assertEqual(post.tags.all()[0].name, "tag1")
        self.assertEqual(post.tags.all()[1].name, "tag2")

Before explaining how to make this work, we'll break down how this test is actually written.

We're inheriting from a TestCase helper that will be written to make testing migrations possible named TestMigrations and defining for this class two attributes that configure the migrations before and after that we want to test. migrate_from is the last migration we expect to be run on machines we want to deploy to and migrate_to is the latest new migration we're testing before deploying.

class TagsTestCase(TestMigrations):

    migrate_from = '0009_previous_migration'
    migrate_to = '0010_migration_being_tested'

Because our test is about a migration, data modifying migrations in particular, we want to do some setup before the migration in question (0010_migration_being_tested) is run. An extra setup method is defined to do that kind of data setup after 0009_previous_migration has run but before 0010_migration_being_tested.

def setUpBeforeMigration(self, apps):
    BlogPost = apps.get_model('blog', 'Post')
    self.post_id = BlogPost.objects.create(
        title = "A test post with tags",
        body = "",
        tags = "tag1 tag2",
    ).id

Once our test runs this setup, we expect the final 0010_migration_being_tested migration to be run. At that time, one or more test_*() methods we define can do the sort of assertions tests would normally do. In this case, we're making sure data was converted to the new schema correctly.

def test_tags_migrated(self):
    BlogPost = self.apps.get_model('blog', 'Post')
    post = BlogPost.objects.get(id=self.post_id)

    self.assertEqual(post.tags.count(), 2)
    self.assertEqual(post.tags.all()[0].name, "tag1")
    self.assertEqual(post.tags.all()[1].name, "tag2")

Here we've fetched a copy of this Post model's after-migration version and confirmed the value we set up in setUpBeforeMigration() was converted to the new structure.

Now, let's look at that TestMigrations base class that makes this possible. First, the pieces from Django we'll need to import to build our migration-aware test cases.

from django.apps import apps
from django.test import TestCase
from django.db.migrations.executor import MigrationExecutor
from django.db import connection

We'll be extending the TestCase class. In order to control migration running, we'll use MigrationExecutor, which needs the database connection to operate on. Migrations are tied pretty intrinsically to Django applications, so we'll be using django.apps.apps and, in particular, get_containing_app_config() to identify the current app our tests are running in.

class TestMigrations(TestCase):

    @property
    def app(self):
        return apps.get_containing_app_config(type(self).__module__).name

    migrate_from = None
    migrate_to = None

We're starting with a few necessary properties.

app is a dynamic property that'll look up and return the name of the current app.
migrate_to will be defined on our own test case subclass as the name of the migration we're testing.
migrate_from is the migration we want to set up test data in, usually the latest migration that's currently being deployed in the project.

def setUp(self):
    assert self.migrate_from and self.migrate_to, \
        "TestCase '{}' must define migrate_from and migrate_to properties".format(type(self).__name__)
    self.migrate_from = [(self.app, self.migrate_from)]
    self.migrate_to = [(self.app, self.migrate_to)]
    executor = MigrationExecutor(connection)
    old_apps = executor.loader.project_state(self.migrate_from).apps

After insisting the test case class had defined migrate_to and migrate_from migrations, we use the internal MigrationExecutor utility to get a state of the applications as of the older of the two migrations.

We'll use old_apps in our setUpBeforeMigration() to work with old versions of the models from this app. First, we'll run our migrations backwards to return to this original migration and then call the setUpBeforeMigration() method.

## Reverse to the original migration
executor.migrate(self.migrate_from)

self.setUpBeforeMigration(old_apps)

Now that we've set up the old state, we simply run the migrations forward again. If the migrations are correct, they should update any test data we created. Of course, we're validating that in our actual tests.

## Run the migration to test
executor = MigrationExecutor(connection)
executor.loader.build_graph()  # reload.
executor.migrate(self.migrate_to)

And finally, we store a current version of the app configuration that our tests can access and define a no-op setUpBeforeMigration()

self.apps = executor.loader.project_state(self.migrate_to).apps

def setUpBeforeMigration(self, apps):
pass

Here's a complete version:

from django.apps import apps
from django.test import TestCase
from django.db.migrations.executor import MigrationExecutor
from django.db import connection


class TestMigrations(TestCase):

    @property
    def app(self):
        return apps.get_containing_app_config(type(self).__module__).name

    migrate_from = None
    migrate_to = None

    def setUp(self):
        assert self.migrate_from and self.migrate_to, \
            "TestCase '{}' must define migrate_from and migrate_to properties".format(type(self).__name__)
        self.migrate_from = [(self.app, self.migrate_from)]
        self.migrate_to = [(self.app, self.migrate_to)]
        executor = MigrationExecutor(connection)
        old_apps = executor.loader.project_state(self.migrate_from).apps

        # Reverse to the original migration
        executor.migrate(self.migrate_from)

        self.setUpBeforeMigration(old_apps)

        # Run the migration to test
        executor = MigrationExecutor(connection)
        executor.loader.build_graph()  # reload.
        executor.migrate(self.migrate_to)

        self.apps = executor.loader.project_state(self.migrate_to).apps

    def setUpBeforeMigration(self, apps):
        pass


class TagsTestCase(TestMigrations):

    migrate_from = '0009_previous_migration'
    migrate_to = '0010_migration_being_tested'

    def setUpBeforeMigration(self, apps):
        BlogPost = apps.get_model('blog', 'Post')
        self.post_id = BlogPost.objects.create(
            title = "A test post with tags",
            body = "",
            tags = "tag1 tag2",
        ).id

    def test_tags_migrated(self):
        BlogPost = self.apps.get_model('blog', 'Post')
        post = BlogPost.objects.get(id=self.post_id)

        self.assertEqual(post.tags.count(), 2)
        self.assertEqual(post.tags.all()[0].name, "tag1")
        self.assertEqual(post.tags.all()[1].name, "tag2")

Migrations are an essential part of any Django project today, as are thorough and reliable tests. Hopefully, you can combine these two essentials now. These techniques may be particularly helpful when maintaining Django apps and libraries distributed to many users. A single project might only have one primary database in production, but distributed libraries have to have predictable, safe migrations your users can depend on. Now you can make sure of that.

Not every migration needs thorough tests! For simple non-data migrations like adding a new null column or adding a new table, you’d only be redoing the work of testing the migration tooling in Django itself, because those migrations are simple enough not to have any special impact on your data. However, this method of testing migrations can be valuable when you have data migrations or schema migrations that might be sensitive to data, such as changing constraints, renaming tables and columns, or building new indexes.