Fixed #23822 -- Added support for serializing model managers in migration

Thanks to Shai Berger, Loïc Bistuer, Simon Charette, Andrew Godwin,
Tim Graham, Carl Meyer, and others for their review and input.

7 files changed, 311 insertions, 16 deletions

diff --git a/tests/custom_managers/models.py b/tests/custom_managers/models.py
index cecfd2c948..d4ca730629 100644
--- a/tests/custom_managers/models.py
+++ b/tests/custom_managers/models.py
@@ -60,9 +60,16 @@ class BaseCustomManager(models.Manager):
     def manager_only(self):
         return self.all()
 
+
 CustomManager = BaseCustomManager.from_queryset(CustomQuerySet)
 
 
+class DeconstructibleCustomManager(BaseCustomManager.from_queryset(CustomQuerySet)):
+
+    def __init__(self, a, b, c=1, d=2):
+        super(DeconstructibleCustomManager, self).__init__(a)
+
+
 class FunPeopleManager(models.Manager):
     def get_queryset(self):
         return super(FunPeopleManager, self).get_queryset().filter(fun=True)
diff --git a/tests/custom_managers/tests.py b/tests/custom_managers/tests.py
index bd18e7d8ea..c115482802 100644
--- a/tests/custom_managers/tests.py
+++ b/tests/custom_managers/tests.py
@@ -1,9 +1,11 @@
 from __future__ import unicode_literals
 
+from django.db import models
 from django.test import TestCase
 from django.utils import six
 
-from .models import (Book, Car, FunPerson, OneToOneRestrictedModel, Person,
+from .models import (Book, Car, CustomManager, CustomQuerySet,
+    DeconstructibleCustomManager, FunPerson, OneToOneRestrictedModel, Person,
     PersonManager, PublishedBookManager, RelatedModel, RestrictedModel)
 
 
@@ -470,6 +472,44 @@ class CustomManagerTests(TestCase):
             ordered=False,
         )
 
+    def test_deconstruct_default(self):
+        mgr = models.Manager()
+        as_manager, mgr_path, qs_path, args, kwargs = mgr.deconstruct()
+        self.assertFalse(as_manager)
+        self.assertEqual(mgr_path, 'django.db.models.manager.Manager')
+        self.assertEqual(args, ())
+        self.assertEqual(kwargs, {})
+
+    def test_deconstruct_as_manager(self):
+        mgr = CustomQuerySet.as_manager()
+        as_manager, mgr_path, qs_path, args, kwargs = mgr.deconstruct()
+        self.assertTrue(as_manager)
+        self.assertEqual(qs_path, 'custom_managers.models.CustomQuerySet')
+
+    def test_deconstruct_from_queryset(self):
+        mgr = DeconstructibleCustomManager('a', 'b')
+        as_manager, mgr_path, qs_path, args, kwargs = mgr.deconstruct()
+        self.assertFalse(as_manager)
+        self.assertEqual(mgr_path, 'custom_managers.models.DeconstructibleCustomManager')
+        self.assertEqual(args, ('a', 'b',))
+        self.assertEqual(kwargs, {})
+
+        mgr = DeconstructibleCustomManager('x', 'y', c=3, d=4)
+        as_manager, mgr_path, qs_path, args, kwargs = mgr.deconstruct()
+        self.assertFalse(as_manager)
+        self.assertEqual(mgr_path, 'custom_managers.models.DeconstructibleCustomManager')
+        self.assertEqual(args, ('x', 'y',))
+        self.assertEqual(kwargs, {'c': 3, 'd': 4})
+
+    def test_deconstruct_from_queryset_failing(self):
+        mgr = CustomManager('arg')
+        msg = ("Could not find manager BaseCustomManagerFromCustomQuerySet in "
+               "django.db.models.manager.\n"
+               "Please note that you need to inherit from managers you "
+               "dynamically generated with 'from_queryset()'.")
+        with self.assertRaisesMessage(ValueError, msg):
+            mgr.deconstruct()
+
 
 class TestCars(TestCase):
 
diff --git a/tests/migrations/models.py b/tests/migrations/models.py
index fc063f5cb0..1e6ab07f75 100644
--- a/tests/migrations/models.py
+++ b/tests/migrations/models.py
@@ -50,3 +50,21 @@ class UnmigratedModel(models.Model):
     if its migrations directory has not been repointed)
     """
     pass
+
+
+class FoodQuerySet(models.query.QuerySet):
+    pass
+
+
+class BaseFoodManager(models.Manager):
+    def __init__(self, a, b, c=1, d=2):
+        super(BaseFoodManager, self).__init__()
+        self.args = (a, b, c, d)
+
+
+class FoodManager(BaseFoodManager.from_queryset(FoodQuerySet)):
+    use_in_migrations = True
+
+
+class NoMigrationFoodManager(BaseFoodManager.from_queryset(FoodQuerySet)):
+    pass
diff --git a/tests/migrations/test_autodetector.py b/tests/migrations/test_autodetector.py
index 9e7511cc45..46c0102280 100644
--- a/tests/migrations/test_autodetector.py
+++ b/tests/migrations/test_autodetector.py
@@ -9,6 +9,8 @@ from django.db.migrations.loader import MigrationLoader
 from django.db import models, connection
 from django.contrib.auth.models import AbstractBaseUser
 
+from .models import FoodManager, FoodQuerySet
+
 
 class DeconstructableObject(object):
     """
@@ -159,6 +161,13 @@ class AutodetectorTests(TestCase):
     other_pony = ModelState("otherapp", "Pony", [
         ("id", models.AutoField(primary_key=True)),
     ])
+    other_pony_food = ModelState("otherapp", "Pony", [
+        ("id", models.AutoField(primary_key=True)),
+    ], managers=[
+        ('food_qs', FoodQuerySet.as_manager()),
+        ('food_mgr', FoodManager('a', 'b')),
+        ('food_mgr_kwargs', FoodManager('x', 'y', 3, 4)),
+    ])
     other_stable = ModelState("otherapp", "Stable", [("id", models.AutoField(primary_key=True))])
     third_thing = ModelState("thirdapp", "Thing", [("id", models.AutoField(primary_key=True))])
     book = ModelState("otherapp", "Book", [
@@ -456,13 +465,15 @@ class AutodetectorTests(TestCase):
         """Tests autodetection of new models."""
         # Make state
         before = self.make_project_state([])
-        after = self.make_project_state([self.author_empty])
+        after = self.make_project_state([self.other_pony_food])
         autodetector = MigrationAutodetector(before, after)
         changes = autodetector._detect_changes()
         # Right number/type of migrations?
-        self.assertNumberMigrations(changes, 'testapp', 1)
-        self.assertOperationTypes(changes, 'testapp', 0, ["CreateModel"])
-        self.assertOperationAttributes(changes, "testapp", 0, 0, name="Author")
+        self.assertNumberMigrations(changes, 'otherapp', 1)
+        self.assertOperationTypes(changes, 'otherapp', 0, ["CreateModel"])
+        self.assertOperationAttributes(changes, "otherapp", 0, 0, name="Pony")
+        self.assertEqual([name for name, mgr in changes['otherapp'][0].operations[0].managers],
+                         ['food_qs', 'food_mgr', 'food_mgr_kwargs'])
 
     def test_old_model(self):
         """Tests deletion of old models."""
@@ -1406,6 +1417,24 @@ class AutodetectorTests(TestCase):
         self.assertOperationAttributes(changes, 'testapp', 0, 1, name="author", order_with_respect_to="book")
         self.assertNotIn("_order", [name for name, field in changes['testapp'][0].operations[0].fields])
 
+    def test_alter_model_managers(self):
+        """
+        Tests that changing the model managers adds a new operation.
+        """
+        # Make state
+        before = self.make_project_state([self.other_pony])
+        after = self.make_project_state([self.other_pony_food])
+        autodetector = MigrationAutodetector(before, after)
+        changes = autodetector._detect_changes()
+        # Right number/type of migrations?
+        self.assertNumberMigrations(changes, 'otherapp', 1)
+        self.assertOperationTypes(changes, 'otherapp', 0, ["AlterModelManagers"])
+        self.assertOperationAttributes(changes, 'otherapp', 0, 0, name="pony")
+        self.assertEqual([name for name, mgr in changes['otherapp'][0].operations[0].managers],
+                         ['food_qs', 'food_mgr', 'food_mgr_kwargs'])
+        self.assertEqual(changes['otherapp'][0].operations[0].managers[1][1].args, ('a', 'b', 1, 2))
+        self.assertEqual(changes['otherapp'][0].operations[0].managers[2][1].args, ('x', 'y', 3, 4))
+
     def test_swappable_first_inheritance(self):
         """Tests that swappable models get their CreateModel first."""
         # Make state
diff --git a/tests/migrations/test_operations.py b/tests/migrations/test_operations.py
index 0378b86f02..bccfb8e3e2 100644
--- a/tests/migrations/test_operations.py
+++ b/tests/migrations/test_operations.py
@@ -17,6 +17,7 @@ from django.db.utils import IntegrityError, DatabaseError
 from django.test import override_settings
 from django.utils import six
 
+from .models import FoodManager, FoodQuerySet
 from .test_base import MigrationTestBase
 
 
@@ -48,7 +49,7 @@ class OperationTestBase(MigrationTestBase):
         return project_state, new_state
 
     def set_up_test_model(self, app_label, second_model=False, third_model=False,
-            related_model=False, mti_model=False, proxy_model=False,
+            related_model=False, mti_model=False, proxy_model=False, manager_model=False,
             unique_together=False, options=False, db_table=None, index_together=False):
         """
         Creates a test model state and database table.
@@ -142,6 +143,18 @@ class OperationTestBase(MigrationTestBase):
                 options={"proxy": True},
                 bases=['%s.Pony' % app_label],
             ))
+        if manager_model:
+            operations.append(migrations.CreateModel(
+                "Food",
+                fields=[
+                    ("id", models.AutoField(primary_key=True)),
+                ],
+                managers=[
+                    ("food_qs", FoodQuerySet.as_manager()),
+                    ("food_mgr", FoodManager("a", "b")),
+                    ("food_mgr_kwargs", FoodManager("x", "y", 3, 4)),
+                ]
+            ))
 
         return self.apply_operations(app_label, ProjectState(), operations)
 
@@ -186,6 +199,10 @@ class OperationTests(OperationTestBase):
         self.assertEqual(definition[0], "CreateModel")
         self.assertEqual(definition[1], [])
         self.assertEqual(sorted(definition[2].keys()), ["fields", "name"])
+        # And default manager not in set
+        operation = migrations.CreateModel("Foo", fields=[], managers=[("objects", models.Manager())])
+        definition = operation.deconstruct()
+        self.assertNotIn('managers', definition[2])
 
     def test_create_model_with_unique_after(self):
         """
@@ -365,6 +382,37 @@ class OperationTests(OperationTestBase):
         self.assertTableNotExists("test_crummo_unmanagedpony")
         self.assertTableExists("test_crummo_pony")
 
+    def test_create_model_managers(self):
+        """
+        Tests that the managers on a model are set.
+        """
+        project_state = self.set_up_test_model("test_cmoma")
+        # Test the state alteration
+        operation = migrations.CreateModel(
+            "Food",
+            fields=[
+                ("id", models.AutoField(primary_key=True)),
+            ],
+            managers=[
+                ("food_qs", FoodQuerySet.as_manager()),
+                ("food_mgr", FoodManager("a", "b")),
+                ("food_mgr_kwargs", FoodManager("x", "y", 3, 4)),
+            ]
+        )
+        self.assertEqual(operation.describe(), "Create model Food")
+        new_state = project_state.clone()
+        operation.state_forwards("test_cmoma", new_state)
+        self.assertIn(("test_cmoma", "food"), new_state.models)
+        managers = new_state.models["test_cmoma", "food"].managers
+        self.assertEqual(managers[0][0], "food_qs")
+        self.assertIsInstance(managers[0][1], models.Manager)
+        self.assertEqual(managers[1][0], "food_mgr")
+        self.assertIsInstance(managers[1][1], FoodManager)
+        self.assertEqual(managers[1][1].args, ("a", "b", 1, 2))
+        self.assertEqual(managers[2][0], "food_mgr_kwargs")
+        self.assertIsInstance(managers[2][1], FoodManager)
+        self.assertEqual(managers[2][1].args, ("x", "y", 3, 4))
+
     def test_delete_model(self):
         """
         Tests the DeleteModel operation.
@@ -1208,6 +1256,61 @@ class OperationTests(OperationTestBase):
         self.assertEqual(definition[1], [])
         self.assertEqual(definition[2], {'name': "Rider", 'order_with_respect_to': "pony"})
 
+    def test_alter_model_managers(self):
+        """
+        Tests that the managers on a model are set.
+        """
+        project_state = self.set_up_test_model("test_almoma")
+        # Test the state alteration
+        operation = migrations.AlterModelManagers(
+            "Pony",
+            managers=[
+                ("food_qs", FoodQuerySet.as_manager()),
+                ("food_mgr", FoodManager("a", "b")),
+                ("food_mgr_kwargs", FoodManager("x", "y", 3, 4)),
+            ]
+        )
+        self.assertEqual(operation.describe(), "Change managers on Pony")
+        managers = project_state.models["test_almoma", "pony"].managers
+        self.assertEqual(managers, [])
+
+        new_state = project_state.clone()
+        operation.state_forwards("test_almoma", new_state)
+        self.assertIn(("test_almoma", "pony"), new_state.models)
+        managers = new_state.models["test_almoma", "pony"].managers
+        self.assertEqual(managers[0][0], "food_qs")
+        self.assertIsInstance(managers[0][1], models.Manager)
+        self.assertEqual(managers[1][0], "food_mgr")
+        self.assertIsInstance(managers[1][1], FoodManager)
+        self.assertEqual(managers[1][1].args, ("a", "b", 1, 2))
+        self.assertEqual(managers[2][0], "food_mgr_kwargs")
+        self.assertIsInstance(managers[2][1], FoodManager)
+        self.assertEqual(managers[2][1].args, ("x", "y", 3, 4))
+
+    def test_alter_model_managers_emptying(self):
+        """
+        Tests that the managers on a model are set.
+        """
+        project_state = self.set_up_test_model("test_almomae", manager_model=True)
+        # Test the state alteration
+        operation = migrations.AlterModelManagers("Food", managers=[])
+        self.assertEqual(operation.describe(), "Change managers on Food")
+        self.assertIn(("test_almomae", "food"), project_state.models)
+        managers = project_state.models["test_almomae", "food"].managers
+        self.assertEqual(managers[0][0], "food_qs")
+        self.assertIsInstance(managers[0][1], models.Manager)
+        self.assertEqual(managers[1][0], "food_mgr")
+        self.assertIsInstance(managers[1][1], FoodManager)
+        self.assertEqual(managers[1][1].args, ("a", "b", 1, 2))
+        self.assertEqual(managers[2][0], "food_mgr_kwargs")
+        self.assertIsInstance(managers[2][1], FoodManager)
+        self.assertEqual(managers[2][1].args, ("x", "y", 3, 4))
+
+        new_state = project_state.clone()
+        operation.state_forwards("test_almomae", new_state)
+        managers = new_state.models["test_almomae", "food"].managers
+        self.assertEqual(managers, [])
+
     def test_alter_fk(self):
         """
         Tests that creating and then altering an FK works correctly
diff --git a/tests/migrations/test_state.py b/tests/migrations/test_state.py
index 65ea6dc1fa..512f75ec15 100644
--- a/tests/migrations/test_state.py
+++ b/tests/migrations/test_state.py
@@ -3,7 +3,8 @@ from django.db import models
 from django.db.migrations.state import ProjectState, ModelState, InvalidBasesError
 from django.test import TestCase
 
-from .models import ModelWithCustomBase
+from .models import (FoodManager, FoodQuerySet, ModelWithCustomBase,
+    NoMigrationFoodManager)
 
 
 class StateTests(TestCase):
@@ -54,11 +55,56 @@ class StateTests(TestCase):
                 verbose_name = "tome"
                 db_table = "test_tome"
 
+        class Food(models.Model):
+
+            food_mgr = FoodManager('a', 'b')
+            food_qs = FoodQuerySet.as_manager()
+            food_no_mgr = NoMigrationFoodManager('x', 'y')
+
+            class Meta:
+                app_label = "migrations"
+                apps = new_apps
+
+        class FoodNoManagers(models.Model):
+
+            class Meta:
+                app_label = "migrations"
+                apps = new_apps
+
+        class FoodNoDefaultManager(models.Model):
+
+            food_no_mgr = NoMigrationFoodManager('x', 'y')
+            food_mgr = FoodManager('a', 'b')
+            food_qs = FoodQuerySet.as_manager()
+
+            class Meta:
+                app_label = "migrations"
+                apps = new_apps
+
+        mgr1 = FoodManager('a', 'b')
+        mgr2 = FoodManager('x', 'y', c=3, d=4)
+
+        class FoodOrderedManagers(models.Model):
+            # The managers on this model should be orderd by their creation
+            # counter and not by the order in model body
+
+            food_no_mgr = NoMigrationFoodManager('x', 'y')
+            food_mgr2 = mgr2
+            food_mgr1 = mgr1
+
+            class Meta:
+                app_label = "migrations"
+                apps = new_apps
+
         project_state = ProjectState.from_apps(new_apps)
         author_state = project_state.models['migrations', 'author']
         author_proxy_state = project_state.models['migrations', 'authorproxy']
         sub_author_state = project_state.models['migrations', 'subauthor']
         book_state = project_state.models['migrations', 'book']
+        food_state = project_state.models['migrations', 'food']
+        food_no_managers_state = project_state.models['migrations', 'foodnomanagers']
+        food_no_default_manager_state = project_state.models['migrations', 'foodnodefaultmanager']
+        food_order_manager_state = project_state.models['migrations', 'foodorderedmanagers']
 
         self.assertEqual(author_state.app_label, "migrations")
         self.assertEqual(author_state.name, "Author")
@@ -89,26 +135,43 @@ class StateTests(TestCase):
         self.assertEqual(len(sub_author_state.fields), 2)
         self.assertEqual(sub_author_state.bases, ("migrations.author", ))
 
+        # The default manager is used in migrations
+        self.assertEqual([name for name, mgr in food_state.managers], ['food_mgr'])
+        self.assertEqual(food_state.managers[0][1].args, ('a', 'b', 1, 2))
+
+        # No explicit managers defined. Migrations will fall back to the default
+        self.assertEqual(food_no_managers_state.managers, [])
+
+        # food_mgr is used in migration but isn't the default mgr, hence add the
+        # default
+        self.assertEqual([name for name, mgr in food_no_default_manager_state.managers],
+                         ['food_no_mgr', 'food_mgr'])
+        self.assertEqual(food_no_default_manager_state.managers[0][1].__class__, models.Manager)
+        self.assertIsInstance(food_no_default_manager_state.managers[1][1], FoodManager)
+
+        self.assertEqual([name for name, mgr in food_order_manager_state.managers],
+                         ['food_mgr1', 'food_mgr2'])
+        self.assertEqual([mgr.args for name, mgr in food_order_manager_state.managers],
+                         [('a', 'b', 1, 2), ('x', 'y', 3, 4)])
+
     def test_render(self):
         """
         Tests rendering a ProjectState into an Apps.
         """
         project_state = ProjectState()
         project_state.add_model_state(ModelState(
-            "migrations",
-            "Tag",
-            [
+            app_label="migrations",
+            name="Tag",
+            fields=[
                 ("id", models.AutoField(primary_key=True)),
                 ("name", models.CharField(max_length=100)),
                 ("hidden", models.BooleanField()),
             ],
-            {},
-            None,
         ))
         project_state.add_model_state(ModelState(
-            "migrations",
-            "SubTag",
-            [
+            app_label="migrations",
+            name="SubTag",
+            fields=[
                 ('tag_ptr', models.OneToOneField(
                     auto_created=True,
                     primary_key=True,
@@ -118,15 +181,40 @@ class StateTests(TestCase):
                 )),
                 ("awesome", models.BooleanField()),
             ],
-            options={},
             bases=("migrations.Tag",),
         ))
 
+        base_mgr = models.Manager()
+        mgr1 = FoodManager('a', 'b')
+        mgr2 = FoodManager('x', 'y', c=3, d=4)
+        project_state.add_model_state(ModelState(
+            app_label="migrations",
+            name="Food",
+            fields=[
+                ("id", models.AutoField(primary_key=True)),
+            ],
+            managers=[
+                # The ordering we really want is objects, mgr1, mgr2
+                ('default', base_mgr),
+                ('food_mgr2', mgr2),
+                ('food_mgr1', mgr1),
+            ]
+        ))
+
         new_apps = project_state.render()
         self.assertEqual(new_apps.get_model("migrations", "Tag")._meta.get_field_by_name("name")[0].max_length, 100)
         self.assertEqual(new_apps.get_model("migrations", "Tag")._meta.get_field_by_name("hidden")[0].null, False)
+
         self.assertEqual(len(new_apps.get_model("migrations", "SubTag")._meta.local_fields), 2)
 
+        Food = new_apps.get_model("migrations", "Food")
+        managers = sorted(Food._meta.managers)
+        self.assertEqual([mgr.name for _, mgr, _ in managers],
+                         ['default', 'food_mgr1', 'food_mgr2'])
+        self.assertEqual([mgr.__class__ for _, mgr, _ in managers],
+                         [models.Manager, FoodManager, FoodManager])
+        self.assertIs(managers[0][1], Food._default_manager)
+
     def test_render_model_inheritance(self):
         class Book(models.Model):
             title = models.CharField(max_length=1000)
diff --git a/tests/migrations/test_writer.py b/tests/migrations/test_writer.py
index 618f8df925..c42008c9d3 100644
--- a/tests/migrations/test_writer.py
+++ b/tests/migrations/test_writer.py
@@ -19,6 +19,7 @@ from django.utils.deconstruct import deconstructible
 from django.utils.translation import ugettext_lazy as _
 from django.utils.timezone import get_default_timezone, utc, FixedOffset
 
+from custom_managers import models as custom_manager_models
 import custom_migration_operations.operations
 import custom_migration_operations.more_operations
 
@@ -351,3 +352,12 @@ class WriterTests(TestCase):
 
         string = MigrationWriter.serialize(models.CharField(default=DeconstructableInstances))[0]
         self.assertEqual(string, "models.CharField(default=migrations.test_writer.DeconstructableInstances)")
+
+    def test_serialize_managers(self):
+        self.assertSerializedEqual(models.Manager())
+        self.assertSerializedResultEqual(
+            custom_manager_models.CustomQuerySet.as_manager(),
+            ('custom_managers.models.CustomQuerySet.as_manager()', {'import custom_managers.models'})
+        )
+        self.assertSerializedEqual(custom_manager_models.DeconstructibleCustomManager('a', 'b'))
+        self.assertSerializedEqual(custom_manager_models.DeconstructibleCustomManager('x', 'y', c=3, d=4))