[1.7.x] Fixed #22496: Data migrations get transactions again!

1 files changed, 37 insertions, 0 deletions

diff --git a/tests/migrations/test_operations.py b/tests/migrations/test_operations.py
index d294bdcd26..32319f373c 100644
--- a/tests/migrations/test_operations.py
+++ b/tests/migrations/test_operations.py
@@ -752,6 +752,43 @@ class OperationTests(MigrationTestBase):
         self.assertEqual(project_state.render().get_model("test_runpython", "Pony").objects.count(), 6)
         self.assertEqual(project_state.render().get_model("test_runpython", "ShetlandPony").objects.count(), 2)
 
+    def test_run_python_atomic(self):
+        """
+        Tests the RunPython operation correctly handles the "atomic" keyword
+        """
+        project_state = self.set_up_test_model("test_runpythonatomic", mti_model=True)
+        def inner_method(models, schema_editor):
+            Pony = models.get_model("test_runpythonatomic", "Pony")
+            Pony.objects.create(pink=1, weight=3.55)
+            raise ValueError("Adrian hates ponies.")
+        atomic_migration = Migration("test", "test_runpythonatomic")
+        atomic_migration.operations = [migrations.RunPython(inner_method)]
+        non_atomic_migration = Migration("test", "test_runpythonatomic")
+        non_atomic_migration.operations = [migrations.RunPython(inner_method, atomic=False)]
+        # If we're a fully-transactional database, both versions should rollback
+        if connection.features.can_rollback_ddl:
+            self.assertEqual(project_state.render().get_model("test_runpythonatomic", "Pony").objects.count(), 0)
+            with self.assertRaises(ValueError):
+                with connection.schema_editor() as editor:
+                    atomic_migration.apply(project_state, editor)
+            self.assertEqual(project_state.render().get_model("test_runpythonatomic", "Pony").objects.count(), 0)
+            with self.assertRaises(ValueError):
+                with connection.schema_editor() as editor:
+                    non_atomic_migration.apply(project_state, editor)
+            self.assertEqual(project_state.render().get_model("test_runpythonatomic", "Pony").objects.count(), 0)
+        # Otherwise, the non-atomic operation should leave a row there
+        else:
+            self.assertEqual(project_state.render().get_model("test_runpythonatomic", "Pony").objects.count(), 0)
+            with self.assertRaises(ValueError):
+                with connection.schema_editor() as editor:
+                    atomic_migration.apply(project_state, editor)
+            self.assertEqual(project_state.render().get_model("test_runpythonatomic", "Pony").objects.count(), 0)
+            with self.assertRaises(ValueError):
+                with connection.schema_editor() as editor:
+                    non_atomic_migration.apply(project_state, editor)
+            self.assertEqual(project_state.render().get_model("test_runpythonatomic", "Pony").objects.count(), 1)
+
+
 
 class MigrateNothingRouter(object):
     """