django-polymorphic/polymorphic/tests/test_utils.py

41 lines
1.5 KiB
Python

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)
field_reprs = [
"<Model2A: id 1, field1 (CharField)>",
"<Model2D: id 2, field1 (CharField), field2 (CharField), field3 (CharField), field4 (CharField)>",
"<Model2B: id 3, field1 (CharField), field2 (CharField)>",
"<Model2B: id 4, field1 (CharField), field2 (CharField)>",
]
for f, f_repr in zip(Model2A.objects.order_by("pk"), field_reprs):
self.assertEqual(qrepr(f), f_repr)