Django Mixin Rollback DB Test Case

If you want to rollback the database inside a setup method, you can use this mixin.

Useful if you want keep the same input data for two query, for example CreateEntity and CreateEntities.

from contextlib import contextmanager

from django.db import transaction


class SetupTestDataRollbackMixin:
    """Mixin providing context manager for automatic database rollback."""

    @contextmanager
    def with_db_rollback(self):
        """
        Context manager that automatically rollbacks database changes.

        Useful for operations like: create, update, delete.
        """
        with transaction.atomic():
            # Create a savepoint
            sid = transaction.savepoint()
            try:
                yield
            finally:
                # Rollback to the savepoint
                transaction.savepoint_rollback(sid)

Usage

from django.test import TestCase

class MyTestCase(TestCase, SetupTestDataRollbackMixin):

    @classmethod
    def setUpTestData(cls):
        cls.location = LocationFactory.create(
            root="/test_root",
            enabled=True,
        )
        cls.location.save()

    def test_create_test_data(self):
        with self.with_db_rollback():
            LocationModel.objects.filter(id=self.location.id).delete()
            self.assertFalse(
                LocationModel.objects.filter(id=self.location.id).exists()
            )

        self.assertTrue(
            LocationModel.objects.filter(id=self.location.id).exists()
        )