Improve reset_polymorphic_ctype() for reliability and test it.

This function can now be safely used on a set of models.
2017-08-01 12:11:42 +02:00
parent e10deeaebd
commit 171d14f369
2 changed files with 68 additions and 0 deletions
@@ -0,0 +1,36 @@
+from unittest import TestCase
+
+from polymorphic.models import PolymorphicTypeUndefined
+from polymorphic.tests import Model2A, Model2B, Model2C, Model2D
+from polymorphic.tests.test_orm import qrepr
+from polymorphic.utils import sort_by_subclass, reset_polymorphic_ctype
+
+
+class UtilsTests(TestCase):
+
+    def test_sort_by_subclass(self):
+        self.assertEqual(
+            sort_by_subclass(Model2D, Model2B, Model2D, Model2A, Model2C),
+            [Model2A, Model2B, Model2C, Model2D, Model2D]
+        )
+
+    def test_reset_polymorphic_ctype(self):
+        """
+        Test the the polymorphic_ctype_id can be restored.
+        """
+        Model2A.objects.create(field1='A1')
+        Model2D.objects.create(field1='A1', field2='B2', field3='C3', field4='D4')
+        Model2B.objects.create(field1='A1', field2='B2')
+        Model2B.objects.create(field1='A1', field2='B2')
+        Model2A.objects.all().update(polymorphic_ctype_id=None)
+
+        with self.assertRaises(PolymorphicTypeUndefined):
+            list(Model2A.objects.all())
+
+        reset_polymorphic_ctype(Model2D, Model2B, Model2D, Model2A, Model2C)
+        self.assertEqual(qrepr(Model2A.objects.all()), (
+            '[ <Model2A: id 1, field1 (CharField)>,\n'
+            '  <Model2D: id 2, field1 (CharField), field2 (CharField), field3 (CharField), field4 (CharField)>,\n'
+            '  <Model2B: id 3, field1 (CharField), field2 (CharField)>,\n'
+            '  <Model2B: id 4, field1 (CharField), field2 (CharField)> ]'
+        ))