Django Data Migration Test Example

django_data_migration_test_example.py

# -----------------------------------------------------------
# data_migration_test_base.py
# -----------------------------------------------------------

from importlib import import_module

from django.test import TestCase


class DataFixTestCaseBase(TestCase):
    app_name = None
    migration_name = None
    data_forward = None

    @classmethod
    def setUpClass(cls):
        super().setUpClass()
        if cls.app_name is None:
            raise Exception("`app_name` is not defined")
        if cls.migration_name is None:
            raise Exception("`migration_name` is not defined")

        migration_path = f"{cls.app_name}.migrations.{cls.migration_name}"
        migration_module = import_module(migration_path)

        if not hasattr(migration_module, "data_forward"):
            raise Exception("`data_forward` function is not defined in migration")

        def _f(self, *args):
            migration_module.data_forward(*args)
        cls.data_forward = _f


# -----------------------------------------------------------
# animals_app/tests/test_migrations_002.py
# -----------------------------------------------------------

from importlib import import_module

from django.test import TestCase

from animals_app.models import Animal


class DataFixTestCase(DataFixTestCaseBase):
    app_name = "animals_app"
    migration_name = "0002_datafix__animal__addsuffixtoname"

    def test_data_fix(self):
        # Prepare
        dog = Animal.objects.create(name="Dog", species="dog")
        cat = Animal.objects.create(name="Cat", species="cat")

        # Action
        self.data_forward(Animal)

        # Assertions
        self.assertEqual(Animal.objects.get(id=dog.id).name, "Dog ZZ")
        self.assertEqual(Animal.objects.get(id=cat.id).name, "Cat ZZ")


# -----------------------------------------------------------
# animals_app/tests/0002_datafix__animal__addsuffixtoname.py
# -----------------------------------------------------------

from django.db import migrations


def data_forward(Animal):
    """Just a sample data migration"""
    for animal in Animal.objects.all():
        animal.name += " ZZ"
        animal.save()


def forward(apps, schema_editor):
    Animal = apps.get_model("animals_app", "Animal")
    data_forward(Animal)


class Migration(migrations.Migration):

    dependencies = [
        ('animals_app', '0001_some_previous_nondata_migration'),
    ]

    operations = [
        migrations.RunPython(forward, migrations.RunPython.noop),
    ]


# -----------------------------------------------------------
# animals_app/models.py
# -----------------------------------------------------------

from django.db import models


class Animal(models.Model):
    species = models.CharField(blank=False, null=False, max_length=50)
    name = models.CharField(blank=False, null=False, max_length=100)

    def __str__(self):
        return f"Animal [name={self.name}, species={self.species}]"