How do I run tests against a Django data migration?

Question

Using the following example from the documentation:

def combine_names(apps, schema_editor):
    Person = apps.get_model("yourappname", "Person")
    for person in Person.objects.all():
        person.name = "%s %s" % (person.first_name, person.last_name)
        person.save()

class Migration(migrations.Migration):    
    dependencies = [
        ('yourappname', '0001_initial'),
    ]    
    operations = [
        migrations.RunPython(combine_names),
    ]

How would I create and run a test against this migration, confirming that the data is migrated correctly?

Answer 1

You can use django-test-migrations package. It is suited for testing: data migrations, schema migrations, and migrations' order.

Here's how it works:

from django_test_migrations.migrator import Migrator

# You can specify any database alias you need:
migrator = Migrator(database='default')

old_state = migrator.before(('main_app', '0002_someitem_is_clean'))
SomeItem = old_state.apps.get_model('main_app', 'SomeItem')

# One instance will be `clean`, the other won't be:
SomeItem.objects.create(string_field='a')
SomeItem.objects.create(string_field='a b')

assert SomeItem.objects.count() == 2
assert SomeItem.objects.filter(is_clean=True).count() == 2

new_state = migrator.after(('main_app', '0003_auto_20191119_2125'))
SomeItem = new_state.apps.get_model('main_app', 'SomeItem')

assert SomeItem.objects.count() == 2
# One instance is clean, the other is not:
assert SomeItem.objects.filter(is_clean=True).count() == 1
assert SomeItem.objects.filter(is_clean=False).count() == 1

We also have native integrations for both pytest:

@pytest.mark.django_db
def test_main_migration0002(migrator):
    """Ensures that the second migration works."""
    old_state = migrator.before(('main_app', '0002_someitem_is_clean'))
    SomeItem = old_state.apps.get_model('main_app', 'SomeItem')
    ...

And unittest:

from django_test_migrations.contrib.unittest_case import MigratorTestCase

class TestDirectMigration(MigratorTestCase):
    """This class is used to test direct migrations."""

    migrate_from = ('main_app', '0002_someitem_is_clean')
    migrate_to = ('main_app', '0003_auto_20191119_2125')

    def prepare(self):
        """Prepare some data before the migration."""
        SomeItem = self.old_state.apps.get_model('main_app', 'SomeItem')
        SomeItem.objects.create(string_field='a')
        SomeItem.objects.create(string_field='a b')

    def test_migration_main0003(self):
        """Run the test itself."""
        SomeItem = self.new_state.apps.get_model('main_app', 'SomeItem')

        assert SomeItem.objects.count() == 2
        assert SomeItem.objects.filter(is_clean=True).count() == 1

• Full guide: https://sobolevn.me/2019/10/testing-django-migrations
• Github: https://github.com/wemake-services/django-test-migrations
• PyPI: https://pypi.org/project/django-test-migrations/

Answer 2

I was doing some google to address the same question and found an article that nailed the hammer on the nail for me and seemed less hacky than existing answers. So, putting this here in case it helps anyone else coming though.

The proposed the following subclass of Django's TestCase:

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

And an example use case that they proposed was:

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")

Answer 3

EDIT:

These other answers make more sense:

• https://stackoverflow.com/a/56212859
• https://stackoverflow.com/a/59016744, if you don't mind the extra (dev) dependency

ORIGINAL:

Running your data-migration functions (such as combine_names from the OP's example) through some basic unit-tests, before actually applying them, makes sense to me too.

At first glance this should not be much more difficult than your normal Django unit-tests: migrations are Python modules and the migrations/ folder is a package, so it is possible to import things from them. However, it took some time to get this working.

The first difficulty arises due to the fact that the default migration file names start with a number. For example, suppose the code from the OP's (i.e. Django's) data-migration example sits in 0002_my_data_migration.py, then it is tempting to use

from yourappname.migrations.0002_my_data_migration import combine_names

but that would raise a SyntaxError because the module name starts with a number (0).

There are at least two ways to make this work:

1. Rename the migration file so it does not start with a number. This should be perfectly fine according to the docs: "Django just cares that each migration has a different name." Then you can just use import as above.

2. If you want to stick to the default numbered migration file names, you can use Python's import_module (see docs and this SO question).

The second difficulty arises from the fact that your data-migration functions are designed to be passed into RunPython (docs), so they expect two input arguments by default: apps and schema_editor. To see where these come from, you can inspect the source.

Now, I'm not sure this works for every case (please, anyone, comment if you can clarify), but for our case, it was sufficient to import apps from django.apps and get the schema_editor from the active database connection (django.db.connection).

The following is a stripped-down example showing how you can implement this for the OP example, assuming the migration file is called 0002_my_data_migration.py:

from importlib import import_module
from django.test import TestCase
from django.apps import apps
from django.db import connection
from yourappname.models import Person
# Our filename starts with a number, so we use import_module
data_migration = import_module('yourappname.migrations.0002_my_data_migration')


class DataMigrationTests(TestCase):
    def __init__(self, *args, **kwargs):
        super(DataMigrationTests, self).__init__(*args, **kwargs)
        # Some test values
        self.first_name = 'John'
        self.last_name = 'Doe'
        
    def test_combine_names(self):
        # Create a dummy Person
        Person.objects.create(first_name=self.first_name,
                              last_name=self.last_name, 
                              name=None)
        # Run the data migration function
        data_migration.combine_names(apps, connection.schema_editor())
        # Test the result
        person = Person.objects.get(id=1)
        self.assertEqual('{} {}'.format(self.first_name, self.last_name), person.name)