Moved migrations.test_operations.OperationTestBase to migrations.test_base.

Co-Authored-By: Daniel Tao <daniel.tao@gmail.com>

3 files changed, 151 insertions, 153 deletions

diff --git a/tests/migrations/test_base.py b/tests/migrations/test_base.py
index 970998f562..7c1cf82e16 100644
--- a/tests/migrations/test_base.py
+++ b/tests/migrations/test_base.py
@@ -5,12 +5,16 @@ from contextlib import contextmanager
 from importlib import import_module
 
 from django.apps import apps
-from django.db import connections
+from django.db import connection, connections, migrations, models
+from django.db.migrations.migration import Migration
 from django.db.migrations.recorder import MigrationRecorder
+from django.db.migrations.state import ProjectState
 from django.test import TransactionTestCase
 from django.test.utils import extend_sys_path
 from django.utils.module_loading import module_dir
 
+from .models import FoodManager, FoodQuerySet
+
 
 class MigrationTestBase(TransactionTestCase):
     """
@@ -130,3 +134,147 @@ class MigrationTestBase(TransactionTestCase):
                 new_module = os.path.basename(target_dir) + '.migrations'
                 with self.settings(MIGRATION_MODULES={app_label: new_module}):
                     yield target_migrations_dir
+
+
+class OperationTestBase(MigrationTestBase):
+    """Common functions to help test operations."""
+
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        cls._initial_table_names = frozenset(connection.introspection.table_names())
+
+    def tearDown(self):
+        self.cleanup_test_tables()
+        super().tearDown()
+
+    def cleanup_test_tables(self):
+        table_names = frozenset(connection.introspection.table_names()) - self._initial_table_names
+        with connection.schema_editor() as editor:
+            with connection.constraint_checks_disabled():
+                for table_name in table_names:
+                    editor.execute(editor.sql_delete_table % {
+                        'table': editor.quote_name(table_name),
+                    })
+
+    def apply_operations(self, app_label, project_state, operations, atomic=True):
+        migration = Migration('name', app_label)
+        migration.operations = operations
+        with connection.schema_editor(atomic=atomic) as editor:
+            return migration.apply(project_state, editor)
+
+    def unapply_operations(self, app_label, project_state, operations, atomic=True):
+        migration = Migration('name', app_label)
+        migration.operations = operations
+        with connection.schema_editor(atomic=atomic) as editor:
+            return migration.unapply(project_state, editor)
+
+    def make_test_state(self, app_label, operation, **kwargs):
+        """
+        Makes a test state using set_up_test_model and returns the
+        original state and the state after the migration is applied.
+        """
+        project_state = self.set_up_test_model(app_label, **kwargs)
+        new_state = project_state.clone()
+        operation.state_forwards(app_label, new_state)
+        return project_state, new_state
+
+    def set_up_test_model(
+        self, app_label, second_model=False, third_model=False, index=False,
+        multicol_index=False, related_model=False, mti_model=False,
+        proxy_model=False, manager_model=False, unique_together=False,
+        options=False, db_table=None, index_together=False, constraints=None,
+    ):
+        """Creates a test model state and database table."""
+        # Make the "current" state.
+        model_options = {
+            'swappable': 'TEST_SWAP_MODEL',
+            'index_together': [['weight', 'pink']] if index_together else [],
+            'unique_together': [['pink', 'weight']] if unique_together else [],
+        }
+        if options:
+            model_options['permissions'] = [('can_groom', 'Can groom')]
+        if db_table:
+            model_options['db_table'] = db_table
+        operations = [migrations.CreateModel(
+            'Pony',
+            [
+                ('id', models.AutoField(primary_key=True)),
+                ('pink', models.IntegerField(default=3)),
+                ('weight', models.FloatField()),
+            ],
+            options=model_options,
+        )]
+        if index:
+            operations.append(migrations.AddIndex(
+                'Pony',
+                models.Index(fields=['pink'], name='pony_pink_idx'),
+            ))
+        if multicol_index:
+            operations.append(migrations.AddIndex(
+                'Pony',
+                models.Index(fields=['pink', 'weight'], name='pony_test_idx'),
+            ))
+        if constraints:
+            for constraint in constraints:
+                operations.append(migrations.AddConstraint('Pony', constraint))
+        if second_model:
+            operations.append(migrations.CreateModel(
+                'Stable',
+                [
+                    ('id', models.AutoField(primary_key=True)),
+                ]
+            ))
+        if third_model:
+            operations.append(migrations.CreateModel(
+                'Van',
+                [
+                    ('id', models.AutoField(primary_key=True)),
+                ]
+            ))
+        if related_model:
+            operations.append(migrations.CreateModel(
+                'Rider',
+                [
+                    ('id', models.AutoField(primary_key=True)),
+                    ('pony', models.ForeignKey('Pony', models.CASCADE)),
+                    ('friend', models.ForeignKey('self', models.CASCADE))
+                ],
+            ))
+        if mti_model:
+            operations.append(migrations.CreateModel(
+                'ShetlandPony',
+                fields=[
+                    ('pony_ptr', models.OneToOneField(
+                        'Pony',
+                        models.CASCADE,
+                        auto_created=True,
+                        parent_link=True,
+                        primary_key=True,
+                        to_field='id',
+                        serialize=False,
+                    )),
+                    ('cuteness', models.IntegerField(default=1)),
+                ],
+                bases=['%s.Pony' % app_label],
+            ))
+        if proxy_model:
+            operations.append(migrations.CreateModel(
+                'ProxyPony',
+                fields=[],
+                options={'proxy': True},
+                bases=['%s.Pony' % app_label],
+            ))
+        if manager_model:
+            operations.append(migrations.CreateModel(
+                'Food',
+                fields=[
+                    ('id', models.AutoField(primary_key=True)),
+                ],
+                managers=[
+                    ('food_qs', FoodQuerySet.as_manager()),
+                    ('food_mgr', FoodManager('a', 'b')),
+                    ('food_mgr_kwargs', FoodManager('x', 'y', 3, 4)),
+                ]
+            ))
+        return self.apply_operations(app_label, ProjectState(), operations)
diff --git a/tests/migrations/test_multidb.py b/tests/migrations/test_multidb.py
index e0c5a4d3c5..d8986deb81 100644
--- a/tests/migrations/test_multidb.py
+++ b/tests/migrations/test_multidb.py
@@ -2,7 +2,7 @@ from django.db import connection, migrations, models
 from django.db.migrations.state import ProjectState
 from django.test import override_settings
 
-from .test_operations import OperationTestBase
+from .test_base import OperationTestBase
 
 
 class AgnosticRouter:
diff --git a/tests/migrations/test_operations.py b/tests/migrations/test_operations.py
index 3c321d4cb4..33b76984d4 100644
--- a/tests/migrations/test_operations.py
+++ b/tests/migrations/test_operations.py
@@ -10,163 +10,13 @@ from django.db.utils import IntegrityError
 from django.test import SimpleTestCase, override_settings, skipUnlessDBFeature
 
 from .models import FoodManager, FoodQuerySet, UnicodeModel
-from .test_base import MigrationTestBase
+from .test_base import OperationTestBase
 
 
 class Mixin:
     pass
 
 
-class OperationTestBase(MigrationTestBase):
-    """
-    Common functions to help test operations.
-    """
-
-    @classmethod
-    def setUpClass(cls):
-        super().setUpClass()
-        cls._initial_table_names = frozenset(connection.introspection.table_names())
-
-    def tearDown(self):
-        self.cleanup_test_tables()
-        super().tearDown()
-
-    def cleanup_test_tables(self):
-        table_names = frozenset(connection.introspection.table_names()) - self._initial_table_names
-        with connection.schema_editor() as editor:
-            with connection.constraint_checks_disabled():
-                for table_name in table_names:
-                    editor.execute(editor.sql_delete_table % {
-                        'table': editor.quote_name(table_name),
-                    })
-
-    def apply_operations(self, app_label, project_state, operations, atomic=True):
-        migration = Migration('name', app_label)
-        migration.operations = operations
-        with connection.schema_editor(atomic=atomic) as editor:
-            return migration.apply(project_state, editor)
-
-    def unapply_operations(self, app_label, project_state, operations, atomic=True):
-        migration = Migration('name', app_label)
-        migration.operations = operations
-        with connection.schema_editor(atomic=atomic) as editor:
-            return migration.unapply(project_state, editor)
-
-    def make_test_state(self, app_label, operation, **kwargs):
-        """
-        Makes a test state using set_up_test_model and returns the
-        original state and the state after the migration is applied.
-        """
-        project_state = self.set_up_test_model(app_label, **kwargs)
-        new_state = project_state.clone()
-        operation.state_forwards(app_label, new_state)
-        return project_state, new_state
-
-    def set_up_test_model(
-            self, app_label, second_model=False, third_model=False, index=False, multicol_index=False,
-            related_model=False, mti_model=False, proxy_model=False, manager_model=False,
-            unique_together=False, options=False, db_table=None, index_together=False, constraints=None):
-        """
-        Creates a test model state and database table.
-        """
-        # Make the "current" state
-        model_options = {
-            "swappable": "TEST_SWAP_MODEL",
-            "index_together": [["weight", "pink"]] if index_together else [],
-            "unique_together": [["pink", "weight"]] if unique_together else [],
-        }
-        if options:
-            model_options["permissions"] = [("can_groom", "Can groom")]
-        if db_table:
-            model_options["db_table"] = db_table
-        operations = [migrations.CreateModel(
-            "Pony",
-            [
-                ("id", models.AutoField(primary_key=True)),
-                ("pink", models.IntegerField(default=3)),
-                ("weight", models.FloatField()),
-            ],
-            options=model_options,
-        )]
-        if index:
-            operations.append(migrations.AddIndex(
-                "Pony",
-                models.Index(fields=["pink"], name="pony_pink_idx")
-            ))
-        if multicol_index:
-            operations.append(migrations.AddIndex(
-                "Pony",
-                models.Index(fields=["pink", "weight"], name="pony_test_idx")
-            ))
-        if constraints:
-            for constraint in constraints:
-                operations.append(migrations.AddConstraint(
-                    "Pony",
-                    constraint,
-                ))
-        if second_model:
-            operations.append(migrations.CreateModel(
-                "Stable",
-                [
-                    ("id", models.AutoField(primary_key=True)),
-                ]
-            ))
-        if third_model:
-            operations.append(migrations.CreateModel(
-                "Van",
-                [
-                    ("id", models.AutoField(primary_key=True)),
-                ]
-            ))
-        if related_model:
-            operations.append(migrations.CreateModel(
-                "Rider",
-                [
-                    ("id", models.AutoField(primary_key=True)),
-                    ("pony", models.ForeignKey("Pony", models.CASCADE)),
-                    ("friend", models.ForeignKey("self", models.CASCADE))
-                ],
-            ))
-        if mti_model:
-            operations.append(migrations.CreateModel(
-                "ShetlandPony",
-                fields=[
-                    ('pony_ptr', models.OneToOneField(
-                        'Pony',
-                        models.CASCADE,
-                        auto_created=True,
-                        parent_link=True,
-                        primary_key=True,
-                        to_field='id',
-                        serialize=False,
-                    )),
-                    ("cuteness", models.IntegerField(default=1)),
-                ],
-                bases=['%s.Pony' % app_label],
-            ))
-        if proxy_model:
-            operations.append(migrations.CreateModel(
-                "ProxyPony",
-                fields=[],
-                options={"proxy": True},
-                bases=['%s.Pony' % app_label],
-            ))
-        if manager_model:
-            operations.append(migrations.CreateModel(
-                "Food",
-                fields=[
-                    ("id", models.AutoField(primary_key=True)),
-                ],
-                managers=[
-                    ("food_qs", FoodQuerySet.as_manager()),
-                    ("food_mgr", FoodManager("a", "b")),
-                    ("food_mgr_kwargs", FoodManager("x", "y", 3, 4)),
-                ]
-            ))
-
-        return self.apply_operations(app_label, ProjectState(), operations)
-
-
 class OperationTests(OperationTestBase):
     """
     Tests running the operations and making sure they do what they say they do.