Added missing test for RunPython behavior when unapplying.

1 files changed, 27 insertions, 4 deletions

diff --git a/tests/migrations/test_operations.py b/tests/migrations/test_operations.py
index 40acd5ed61..86dbb84411 100644
--- a/tests/migrations/test_operations.py
+++ b/tests/migrations/test_operations.py
@@ -2009,10 +2009,11 @@ class OperationTests(OperationTestBase):
             Pony.objects.create(pink=1, weight=3.55)
             raise ValueError("Adrian hates ponies.")
 
+        # Verify atomicity when applying.
         atomic_migration = Migration("test", "test_runpythonatomic")
-        atomic_migration.operations = [migrations.RunPython(inner_method)]
+        atomic_migration.operations = [migrations.RunPython(inner_method, reverse_code=inner_method)]
         non_atomic_migration = Migration("test", "test_runpythonatomic")
-        non_atomic_migration.operations = [migrations.RunPython(inner_method, atomic=False)]
+        non_atomic_migration.operations = [migrations.RunPython(inner_method, reverse_code=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.apps.get_model("test_runpythonatomic", "Pony").objects.count(), 0)
@@ -2035,11 +2036,33 @@ class OperationTests(OperationTestBase):
                 with connection.schema_editor() as editor:
                     non_atomic_migration.apply(project_state, editor)
             self.assertEqual(project_state.apps.get_model("test_runpythonatomic", "Pony").objects.count(), 1)
-        # And deconstruction
+        # Reset object count to zero and verify atomicity when unapplying.
+        project_state.apps.get_model("test_runpythonatomic", "Pony").objects.all().delete()
+        # On a fully-transactional database, both versions rollback.
+        if connection.features.can_rollback_ddl:
+            self.assertEqual(project_state.apps.get_model("test_runpythonatomic", "Pony").objects.count(), 0)
+            with self.assertRaises(ValueError):
+                with connection.schema_editor() as editor:
+                    atomic_migration.unapply(project_state, editor)
+            self.assertEqual(project_state.apps.get_model("test_runpythonatomic", "Pony").objects.count(), 0)
+            with self.assertRaises(ValueError):
+                with connection.schema_editor() as editor:
+                    non_atomic_migration.unapply(project_state, editor)
+            self.assertEqual(project_state.apps.get_model("test_runpythonatomic", "Pony").objects.count(), 0)
+        # Otherwise, the non-atomic operation leaves a row there.
+        else:
+            self.assertEqual(project_state.apps.get_model("test_runpythonatomic", "Pony").objects.count(), 0)
+            with self.assertRaises(ValueError):
+                with connection.schema_editor() as editor:
+                    atomic_migration.unapply(project_state, editor)
+            with self.assertRaises(ValueError):
+                with connection.schema_editor() as editor:
+                    non_atomic_migration.unapply(project_state, editor)
+        # Verify deconstruction.
         definition = non_atomic_migration.operations[0].deconstruct()
         self.assertEqual(definition[0], "RunPython")
         self.assertEqual(definition[1], [])
-        self.assertEqual(sorted(definition[2]), ["atomic", "code"])
+        self.assertEqual(sorted(definition[2]), ["atomic", "code", "reverse_code"])
 
     def test_run_python_related_assignment(self):
         """