Caktus Consulting Group: Writing Unit Tests for Django Migrations

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:

classTagsTestCase(TestMigrations):migrate_from='0009_previous_migration'migrate_to='0010_migration_being_tested'defsetUpBeforeMigration(self,apps):BlogPost=apps.get_model('blog','Post')self.post_id=BlogPost.objects.create(title="A test post with tags",body="",tags="tag1 tag2",).iddeftest_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.

classTagsTestCase(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 after0009_previous_migration has run but before 0010_migration_being_tested.

defsetUpBeforeMigration(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.

deftest_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.

fromdjango.appsimportappsfromdjango.testimportTransactionTestCasefromdjango.db.migrations.executorimportMigrationExecutorfromdjango.dbimportconnection

We'll be extending the TransactionTestCase 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.

classTestMigrations(TransactionTestCase):@propertydefapp(self):returnapps.get_containing_app_config(type(self).__module__).namemigrate_from=Nonemigrate_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 been deployed in the project.

defsetUp(self):assertself.migrate_fromandself.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 migrationexecutor.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 testexecutor.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).appsdefsetUpBeforeMigration(self,apps):pass

Here's a complete version:

fromdjango.appsimportappsfromdjango.testimportTransactionTestCasefromdjango.db.migrations.executorimportMigrationExecutorfromdjango.dbimportconnectionclassTestMigrations(TransactionTestCase):@propertydefapp(self):returnapps.get_containing_app_config(type(self).__module__).namemigrate_from=Nonemigrate_to=NonedefsetUp(self):assertself.migrate_fromandself.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 migrationexecutor.migrate(self.migrate_from)self.setUpBeforeMigration(old_apps)# Run the migration to testexecutor.migrate(self.migrate_to)self.apps=executor.loader.project_state(self.migrate_to).appsdefsetUpBeforeMigration(self,apps):passclassTagsTestCase(TestMigrations):migrate_from='0009_previous_migration'migrate_to='0010_migration_being_tested'defsetUpBeforeMigration(self,apps):BlogPost=apps.get_model('blog','Post')self.post_id=BlogPost.objects.create(title="A test post with tags",body="",tags="tag1 tag2",).iddeftest_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")