A pytest fixture to test Django data migrations

conftest.py

# based on https://gist.github.com/blueyed/4fb0a807104551f103e6
# and on https://gist.github.com/TauPan/aec52e398d7288cb5a62895916182a9f  (gistspection!)

from django.core.management import call_command
from django.db import connection
from django.db.migrations.executor import MigrationExecutor

import pytest


@pytest.fixture()
def migration(request, transactional_db):
    # see https://gist.github.com/asfaltboy/b3e6f9b5d95af8ba2cc46f2ba6eae5e2
    """
    This fixture returns a helper object to test Django data migrations.
    The fixture returns an object with two methods;
     - `before` to initialize db to the state before the migration under test
     - `after` to execute the migration and bring db to the state after the migration
    The methods return `old_apps` and `new_apps` respectively; these can
    be used to initiate the ORM models as in the migrations themselves.
    For example:

        def test_foo_set_to_bar(migration):
            old_apps = migration.before([('my_app', '0001_inital')])

            Foo = old_apps.get_model('my_app', 'foo')
            Foo.objects.create(bar=False)

            assert Foo.objects.count() == 1
            assert Foo.objects.filter(bar=False).count() == Foo.objects.count()

            # executing migration
            new_apps = migration.apply([('my_app', '0002_set_foo_bar')])
            Foo = new_apps.get_model('my_app', 'foo')
 
            assert Foo.objects.filter(bar=False).count() == 0
            assert Foo.objects.filter(bar=True).count() == Foo.objects.count()

    Based on: https://gist.github.com/blueyed/4fb0a807104551f103e6
    """
    class Migrator(object):
        def before(self, targets):
            """ Specify app and starting migration names as in:
                before([('app', '0001_before')]) => app/migrations/0001_before.py
            """
            self.executor = MigrationExecutor(connection)
            # prepare state of db to before the migration ("migrate_from" state)
            self._old_apps = self.executor.migrate(targets).apps
            return self._old_apps

        def apply(self, targets):
            """ Migrate forwards to the "targets" migration """
            self.executor.loader.build_graph()  # reload.
            self._new_apps = self.executor.migrate(targets).apps
            return self._new_apps

    # ensure to migrate forward migrated apps all the way after test
    def migrate_to_end():
        call_command('migrate', verbosity=0)
    request.addfinalizer(migrate_to_end)

    return Migrator()

I have made a pypi package out of this gist with some extra features.
Check it out: https://github.com/wemake-services/django-test-migrations