Allow callables as the argument to RunPython

1 files changed, 9 insertions, 0 deletions

diff --git a/tests/migrations/test_operations.py b/tests/migrations/test_operations.py
index 1f247c1a97..db79e730e2 100644
--- a/tests/migrations/test_operations.py
+++ b/tests/migrations/test_operations.py
@@ -332,6 +332,15 @@ class OperationTests(MigrationTestBase):
         # 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):
+            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)
+        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)
 
 
 class MigrateNothingRouter(object):