fix and refactor migrations test teardown and setup

django_test_migrations/contrib/pytest_plugin.py

@@ -5,7 +5,7 @@
 
 
 @pytest.fixture()
-def migrator_factory(transactional_db, django_db_use_migrations):
+def migrator_factory(request, transactional_db, django_db_use_migrations):
     """
     Pytest fixture to create migrators inside the pytest tests.
 
@@ -36,9 +36,9 @@ def test_migration(migrator_factory):
         pytest.skip('--nomigrations was specified')
 
     def factory(database_name: Optional[str] = None) -> Migrator:
-        # ``Migrator.reset`` is not registered as finalizer here, because
-        # database is flushed by ``transactional_db`` fixture's finalizers
-        return Migrator(database_name)
+        migrator = Migrator(database_name)
+        request.addfinalizer(migrator.reset)
+        return migrator
     return factory
 
 

django_test_migrations/contrib/unittest_case.py

@@ -3,7 +3,8 @@
 from django.db.migrations.state import ProjectState
 from django.test import TransactionTestCase
 
-from django_test_migrations.migrator import Migrator, _Migration  # noqa: WPS450
+from django_test_migrations.migrator import Migrator
+from django_test_migrations.types import MigrationSpec
 
 
 class MigratorTestCase(TransactionTestCase):
@@ -14,8 +15,8 @@ class MigratorTestCase(TransactionTestCase):
     new_state: ProjectState
 
     #: Part of the end-user API. Used to tell what migrations we are using.
-    migrate_from: ClassVar[_Migration]
-    migrate_to: ClassVar[_Migration]
+    migrate_from: ClassVar[MigrationSpec]
+    migrate_to: ClassVar[MigrationSpec]
 
     def setUp(self) -> None:
         """
@@ -40,3 +41,8 @@ def prepare(self) -> None:
 
         Used to prepare some data before the migration process.
         """
+
+    def tearDown(self) -> None:
+        """Used to clean mess up after each test."""
+        self._migrator.reset()
+        super().tearDown()

django_test_migrations/exceptions.py

@@ -0,0 +1,14 @@
+from django_test_migrations.types import MigrationTarget
+
+
+class MigrationNotInPlan(Exception):
+    """``MigrationTarget`` not found in migrations plan."""
+
+    def __init__(self, migration_target: MigrationTarget) -> None:  # noqa: D107
+        self.migration_target = migration_target
+
+    def __str__(self) -> str:
+        """String representation of exception's instance."""
+        return 'Migration {0} not found in migrations plan'.format(
+            self.migration_target,
+        )

django_test_migrations/logic/migrations.py

@@ -0,0 +1,10 @@
+from typing import List
+
+from django_test_migrations.types import MigrationSpec, MigrationTarget
+
+
+def normalize(migration_target: MigrationSpec) -> List[MigrationTarget]:
+    """Normalize ``migration_target`` to expected format."""
+    if not isinstance(migration_target, list):
+        migration_target = [migration_target]
+    return migration_target

django_test_migrations/migrator.py

@@ -1,16 +1,17 @@
 from contextlib import contextmanager
-from typing import List, Optional, Tuple, Union
+from typing import Optional
 
 from django.core.management import call_command
+from django.core.management.color import no_style
 from django.db import DEFAULT_DB_ALIAS, connections
 from django.db.migrations.executor import MigrationExecutor
 from django.db.migrations.state import ProjectState
 from django.db.models.signals import post_migrate, pre_migrate
 
-# Regular or rollback migration: 0001 -> 0002, or 0002 -> 0001
-# Rollback migration to initial state: 0001 -> None
-_Migration = Tuple[str, Optional[str]]
-_MigrationSpec = Union[_Migration, List[_Migration]]
+from django_test_migrations import sql
+from django_test_migrations.logic.migrations import normalize
+from django_test_migrations.plan import truncate_plan
+from django_test_migrations.types import MigrationPlan, MigrationSpec
 
 
 @contextmanager
@@ -61,23 +62,41 @@ def __init__(
         self._database: str = database
         self._executor = MigrationExecutor(connections[self._database])
 
-    def before(self, migrate_from: _MigrationSpec) -> ProjectState:
+    def before(self, migrate_from: MigrationSpec) -> ProjectState:
         """Reverse back to the original migration."""
-        if not isinstance(migrate_from, list):
-            migrate_from = [migrate_from]
-        with _mute_migrate_signals():
-            return self._executor.migrate(migrate_from)
+        migrate_from = normalize(migrate_from)
+
+        style = no_style()
+        # start from clean database state
+        sql.drop_models_tables(self._database, style)
+        sql.flush_django_migrations_table(self._database, style)
+
+        # prepare as broad plan as possible based on full plan
+        self._executor.loader.build_graph()  # reload
+        full_plan = self._executor.migration_plan(
+            self._executor.loader.graph.leaf_nodes(),
+            clean_start=True,
+        )
+        plan = truncate_plan(migrate_from, full_plan)
 
-    def after(self, migrate_to: _MigrationSpec) -> ProjectState:
+        # apply all migrations from generated plan on clean database
+        # (only forward, so any unexpected migration won't be applied)
+        # to restore database state before tested migration
+        return self._migrate(migrate_from, plan=plan)
+
+    def after(self, migrate_to: MigrationSpec) -> ProjectState:
         """Apply the next migration."""
-        self._executor.loader.build_graph()  # reload.
-        return self.before(migrate_to)
+        self._executor.loader.build_graph()  # reload
+        return self._migrate(normalize(migrate_to))
 
     def reset(self) -> None:
         """Reset the state to the most recent one."""
-        call_command(
-            'flush',
-            database=self._database,
-            interactive=False,
-            verbosity=0,
-        )
+        call_command('migrate', verbosity=0, database=self._database)
+
+    def _migrate(
+        self,
+        migration_targets: MigrationSpec,
+        plan: Optional[MigrationPlan] = None,
+    ) -> ProjectState:
+        with _mute_migrate_signals():
+            return self._executor.migrate(migration_targets, plan=plan)

django_test_migrations/plan.py

@@ -2,9 +2,13 @@
 
 from django.apps import apps
 from django.db import DEFAULT_DB_ALIAS, connections
+from django.db.migrations import Migration
 from django.db.migrations.graph import Node
 from django.db.migrations.loader import MigrationLoader
 
+from django_test_migrations.exceptions import MigrationNotInPlan
+from django_test_migrations.types import MigrationPlan, MigrationTarget
+
 
 def all_migrations(
     database: str = DEFAULT_DB_ALIAS,
@@ -80,3 +84,55 @@ def _generate_plan(
                 plan.append(node)
                 seen.add(migration)
     return plan
+
+
+def truncate_plan(
+    targets: List[MigrationTarget],
+    plan: MigrationPlan,
+) -> MigrationPlan:
+    """Truncate migrations ``plan`` up to ``targets``.
+
+    This method is used mainly to truncate full/clean migrations plan
+    to get as broad plan as possible.
+    By "broad" plan we mean plan with all targets migrations included
+    as well as all older migrations not related with targets.
+    "Broad" plan is needed mostly to make ``Migrator`` API developers'
+    friendly, just to not force developers to include migrations targets
+    for all other models they want to use in test (e.g. to setup some
+    model instances) in ``migrate_from``.
+    Such plan will also produce database state really similar to state
+    from our production environment just before new migrations are applied.
+    Migrations plan for targets generated by Django's
+    ``MigrationExecutor.migration_plan`` is minimum plan needed to apply
+    targets migrations, it includes only migrations targets with all its
+    dependencies, so it doesn't fit to our approach, that's why following
+    function is needed.
+
+    """
+    if not targets or not plan:
+        return plan
+
+    target_max_index = max(_get_index(target, plan) for target in targets)
+    return plan[:(target_max_index + 1)]
+
+
+def _get_index(target: MigrationTarget, plan: MigrationPlan) -> int:
+    try:
+        index = next(
+            index
+            for index, (migration, _) in enumerate(plan)  # noqa: WPS405, WPS414
+            if _filter_predicate(target, migration)
+        )
+    except StopIteration:
+        raise MigrationNotInPlan(target)
+    else:
+        # exclude target app from migrations plan
+        return index - (target[1] is None)
+
+
+def _filter_predicate(target: MigrationTarget, migration: Migration) -> bool:
+    # when ``None`` passed as migration name then initial migration from
+    # target's app will be chosen and handled properly in ``_get_index``
+    # so in final all target app migrations will be excluded from plan
+    index = 2 - (target[1] is None)
+    return (migration.app_label, migration.name)[:index] == target[:index]

django_test_migrations/sql.py

@@ -0,0 +1,113 @@
+from functools import partial
+from typing import Callable, Dict, List, Optional, Union
+
+from django.core.management.color import Style, no_style
+from django.db import DefaultConnectionProxy, connections, transaction
+from django.db.backends.base.base import BaseDatabaseWrapper
+from typing_extensions import Final
+
+_Connection = Union[DefaultConnectionProxy, BaseDatabaseWrapper]
+
+DJANGO_MIGRATIONS_TABLE_NAME: Final = 'django_migrations'
+
+
+def drop_models_tables(
+    database_name: str,
+    style: Optional[Style] = None,
+) -> None:
+    """Drop all installed Django's models tables."""
+    style = style or no_style()
+    connection = connections[database_name]
+    tables = connection.introspection.django_table_names(
+        only_existing=True,
+        include_views=False,
+    )
+    sql_drop_tables = [
+        connection.SchemaEditorClass.sql_delete_table % {
+            'table': style.SQL_FIELD(connection.ops.quote_name(table)),
+        }
+        for table in tables
+    ]
+    if sql_drop_tables:
+        get_execute_sql_flush_for(connection)(database_name, sql_drop_tables)
+
+
+def flush_django_migrations_table(
+    database_name: str,
+    style: Optional[Style] = None,
+) -> None:
+    """Flush `django_migrations` table.
+
+    Ensures compability with all supported Django versions.
+    `django_migrations` is not "regular" Django model, so its not returned
+    by ``ConnectionRouter.get_migratable_models`` which is used e.g. to
+    implement sequences reset in ``Django==1.11``.
+
+    """
+    style = style or no_style()
+    connection = connections[database_name]
+    django_migrations_sequences = get_django_migrations_table_sequences(
+        connection,
+    )
+    execute_sql_flush = get_execute_sql_flush_for(connection)
+    execute_sql_flush(
+        database_name,
+        connection.ops.sql_flush(
+            style,
+            [DJANGO_MIGRATIONS_TABLE_NAME],
+            django_migrations_sequences,
+            allow_cascade=False,
+        ),
+    )
+
+
+def get_django_migrations_table_sequences(
+    connection: _Connection,
+) -> List[Dict[str, str]]:
+    """`django_migrations` table introspected sequences.
+
+    Returns properly inspected sequences when using ``Django>1.11``
+    and static sequence for `id` column otherwise.
+
+    """
+    if hasattr(connection.introspection, 'get_sequences'):  # noqa: WPS421
+        with connection.cursor() as cursor:
+            return connection.introspection.get_sequences(
+                cursor,
+                DJANGO_MIGRATIONS_TABLE_NAME,
+            )
+    # for ``Django==1.11`` only primary key sequence is returned
+    return [{'table': DJANGO_MIGRATIONS_TABLE_NAME, 'column': 'id'}]
+
+
+def get_execute_sql_flush_for(
+    connection: _Connection,
+) -> Callable[[str, List[str]], None]:
+    """Return ``execute_sql_flush`` callable for given connection."""
+    return getattr(
+        connection.ops,
+        'execute_sql_flush',
+        partial(execute_sql_flush, connection),
+    )
+
+
+def execute_sql_flush(
+    connection: _Connection,
+    using: str,
+    sql_list: List[str],
+) -> None:  # pragma: no cover
+    """Execute a list of SQL statements to flush the database.
+
+    This function is copy of ``connection.ops.execute_sql_flush``
+    method from Django's source code: https://bit.ly/3doGMye
+    to make `django-test-migrations` compatible with ``Django==1.11``.
+    ``connection.ops.execute_sql_flush()`` was introduced in ``Django==2.0``.
+
+    """
+    with transaction.atomic(
+        using=using,
+        savepoint=connection.features.can_rollback_ddl,
+    ):
+        with connection.cursor() as cursor:
+            for sql in sql_list:
+                cursor.execute(sql)

django_test_migrations/types.py

@@ -0,0 +1,11 @@
+from typing import List, Optional, Tuple, Union
+
+from django.db.migrations import Migration
+
+# Migration target: (app_name, migration_name)
+# Regular or rollback migration: 0001 -> 0002, or 0002 -> 0001
+# Rollback migration to initial state: 0001 -> None
+MigrationTarget = Tuple[str, Optional[str]]
+MigrationSpec = Union[MigrationTarget, List[MigrationTarget]]
+
+MigrationPlan = List[Tuple[Migration, bool]]

pyproject.toml

@@ -69,3 +69,4 @@ pytest-cov = "^2.7"
 pytest-randomly = "^3.3"
 pytest-django = "^3.9"
 pytest-pythonpath = "^0.7"
+pytest-mock = "^3.1.0"