Add RunPython migration operation and tests

1 files changed, 28 insertions, 1 deletions

diff --git a/tests/migrations/test_operations.py b/tests/migrations/test_operations.py
index 72f938fe4e..1f247c1a97 100644
--- a/tests/migrations/test_operations.py
+++ b/tests/migrations/test_operations.py
@@ -282,7 +282,7 @@ class OperationTests(MigrationTestBase):
 
     def test_run_sql(self):
         """
-        Tests the AlterIndexTogether operation.
+        Tests the RunSQL operation.
         """
         project_state = self.set_up_test_model("test_runsql")
         # Create the operation
@@ -306,6 +306,33 @@ class OperationTests(MigrationTestBase):
             operation.database_backwards("test_runsql", editor, new_state, project_state)
         self.assertTableNotExists("i_love_ponies")
 
+    def test_run_python(self):
+        """
+        Tests the RunPython operation
+        """
+
+        project_state = self.set_up_test_model("test_runpython")
+        # Create the operation
+        operation = migrations.RunPython(
+            """
+            Pony = models.get_model("test_runpython", "Pony")
+            Pony.objects.create(pink=2, weight=4.55)
+            Pony.objects.create(weight=1)
+            """,
+        )
+        # Test the state alteration does nothing
+        new_state = project_state.clone()
+        operation.state_forwards("test_runpython", new_state)
+        self.assertEqual(new_state, project_state)
+        # Test the database alteration
+        self.assertEqual(project_state.render().get_model("test_runpython", "Pony").objects.count(), 0)
+        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(), 2)
+        # And test reversal fails
+        with self.assertRaises(NotImplementedError):
+            operation.database_backwards("test_runpython", None, new_state, project_state)
+
 
 class MigrateNothingRouter(object):
     """