Ticket #21893: test.diff

tests/migrations/test_operations.py

@@ -14 +14 @@
     both forwards and backwards.
     """
 
-    def set_up_test_model(self, app_label, second_model=False, related_model=False):
+    def set_up_test_model(self, app_label, second_model=False, related_model=False, mti_model=False):
         """
         Creates a test model state and database table.
         """
@@ -38 +38 @@
             ],
         )]
         if second_model:
-            operations.append(migrations.CreateModel("Stable", [("id", models.AutoField(primary_key=True))]))
+            operations.append(migrations.CreateModel(
+                "Stable",
+                [
+                    ("id", models.AutoField(primary_key=True)),
+                ]
+            ))
         if related_model:
             operations.append(migrations.CreateModel(
                 "Rider",
@@ -47 +52 @@
                     ("pony", models.ForeignKey("Pony")),
                 ],
             ))
+        if mti_model:
+            operations.append(migrations.CreateModel(
+                "ShetlandPony",
+                [
+                    ('pony_ptr', models.OneToOneField(
+                        auto_created=True,
+                        primary_key=True,
+                        to_field='id',
+                        serialize=False,
+                        to='Pony',
+                    )),
+                    ("cuteness", models.IntegerField(default=1)),
+                ],
+            ))
         project_state = ProjectState()
         for operation in operations:
             operation.state_forwards(app_label, project_state)
@@ -476 +495 @@
         Tests the RunPython operation
         """
 
-        project_state = self.set_up_test_model("test_runpython")
+        project_state = self.set_up_test_model("test_runpython", mti_model=True)
+
         # Create the operation
         operation = migrations.RunPython(
             """
@@ -497 +517 @@
         # And test reversal fails
         with self.assertRaises(NotImplementedError):
             operation.database_backwards("test_runpython", None, new_state, project_state)
-        # Now test we can do it with a callable
 
-        def inner_method(models, schema_editor):
+        def create_ponies(models, schema_editor):
             Pony = models.get_model("test_runpython", "Pony")
-            Pony.objects.create(pink=1, weight=3.55)
-            Pony.objects.create(weight=5)
-        operation = migrations.RunPython(inner_method)
+            pony1 = Pony.objects.create(pink=1, weight=3.55)
+            self.assertIsNot(pony1.pk, None)
+            pony2 = Pony.objects.create(weight=5)
+            self.assertIsNot(pony2.pk, None)
+            self.assertNotEqual(pony1.pk, pony2.pk)
+
+        operation = migrations.RunPython(create_ponies)
         with connection.schema_editor() as editor:
             operation.database_forwards("test_runpython", editor, project_state, new_state)
         self.assertEqual(project_state.render().get_model("test_runpython", "Pony").objects.count(), 4)
 
+        def create_shetlandponies(models, schema_editor):
+            ShetlandPony = models.get_model("test_runpython", "ShetlandPony")
+            pony1 = ShetlandPony.objects.create()
+            self.assertIsNot(pony1.pk, None)
+            pony2 = ShetlandPony.objects.create(weight=5)
+            self.assertIsNot(pony2.pk, None)
+            self.assertNotEqual(pony1.pk, pony2.pk)
+
+        operation = migrations.RunPython(create_shetlandponies)
+        with connection.schema_editor() as editor:
+            operation.database_forwards("test_runpython", editor, project_state, new_state)
+        self.assertEqual(project_state.render().get_model("test_runpython", "ShetlandPony").objects.count(), 6)
 
 class MigrateNothingRouter(object):
     """