diff --git a/seal/managers.py b/seal/managers.py index 50dba70..1589a63 100644 --- a/seal/managers.py +++ b/seal/managers.py @@ -8,7 +8,7 @@ class SealedModelIterable(models.query.ModelIterable): def __iter__(self): objs = super(SealedModelIterable, self).__iter__() for obj in objs: - obj._state.sealed = True + obj.seal() yield obj diff --git a/seal/models.py b/seal/models.py index b9f113b..052c520 100644 --- a/seal/models.py +++ b/seal/models.py @@ -33,6 +33,9 @@ def __new__(cls, name, bases, attrs): class SealableModel(with_metaclass(SealaleModelBase, models.Model)): objects = SealableQuerySet.as_manager() + def seal(self): + self._state.sealed = True + class Meta: abstract = True diff --git a/tests/models.py b/tests/models.py index 86a05f2..ebf93e1 100644 --- a/tests/models.py +++ b/tests/models.py @@ -1,4 +1,5 @@ from django.db import models +from seal.managers import SealableQuerySet from seal.models import SealableModel @@ -17,3 +18,10 @@ class SeaLion(SealableModel): class GreatSeaLion(SeaLion): # TODO: add support for auto-generated o2os parent_link and non-parent link o2o. sealion_ptr = models.OneToOneField(SeaLion, models.CASCADE, parent_link=True, primary_key=True) + + +class Koala(models.Model): + height = models.PositiveIntegerField() + weight = models.PositiveIntegerField() + + objects = SealableQuerySet.as_manager() diff --git a/tests/test_managers.py b/tests/test_managers.py index d200f1a..3180468 100644 --- a/tests/test_managers.py +++ b/tests/test_managers.py @@ -1,7 +1,7 @@ from django.test import TestCase from seal.exceptions import SealedObject -from .models import GreatSeaLion, Location, SeaLion +from .models import GreatSeaLion, Koala, Location, SeaLion class SealableQuerySetTests(TestCase): @@ -9,6 +9,7 @@ class SealableQuerySetTests(TestCase): def setUpTestData(cls): cls.location = Location.objects.create(latitude=51.585474, longitude=156.634331) cls.great_sealion = GreatSeaLion.objects.create(height=1, weight=100, location=cls.location) + cls.koala = Koala.objects.create(height=1, weight=10) cls.sealion = cls.great_sealion.sealion_ptr cls.sealion.previous_locations.add(cls.location) @@ -106,3 +107,7 @@ def test_not_reverse_sealed_many_to_many(self): def test_sealed_prefetched_reverse_many_to_many(self): instance = Location.objects.prefetch_related('previous_visitors').seal().get() self.assertSequenceEqual(instance.previous_visitors.all(), [self.sealion]) + + def test_improper_usage_raises_error(self): + with self.assertRaises(AttributeError): + Koala.objects.only('pk').seal().get() diff --git a/tests/test_models.py b/tests/test_models.py index 92929fc..fa56cfc 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -7,49 +7,49 @@ class SealableModelTests(SimpleTestCase): def test_sealed_instance_deferred_attribute_access(self): instance = SeaLion.from_db('default', ['id'], [1]) - instance._state.sealed = True + instance.seal() message = "Cannot fetch deferred fields weight on a sealed object." with self.assertRaisesMessage(SealedObject, message): instance.weight def test_sealed_instance_foreign_key_access(self): instance = SeaLion.from_db('default', ['id', 'location_id'], [1, 1]) - instance._state.sealed = True + instance.seal() message = "Cannot fetch related field location on a sealed object." with self.assertRaisesMessage(SealedObject, message): instance.location def test_sealed_instance_reverse_foreign_key_access(self): instance = Location.from_db('default', ['id'], [1]) - instance._state.sealed = True + instance.seal() message = "Cannot fetch many-to-many field visitors on a sealed object." with self.assertRaisesMessage(SealedObject, message): instance.visitors.all() def test_sealed_instance_parent_link_access(self): instance = SeaLion.from_db('default', ['id'], [1]) - instance._state.sealed = True + instance.seal() message = "Cannot fetch related field greatsealion on a sealed object." with self.assertRaisesMessage(SealedObject, message): instance.greatsealion def test_sealed_instance_reverse_parent_link_access(self): instance = GreatSeaLion.from_db('default', ['sealion_ptr_id'], [1]) - instance._state.sealed = True + instance.seal() message = "Cannot fetch related field sealion_ptr on a sealed object." with self.assertRaisesMessage(SealedObject, message): instance.sealion_ptr def test_sealed_instance_m2m_access(self): instance = SeaLion.from_db('default', ['id'], [1]) - instance._state.sealed = True + instance.seal() message = "Cannot fetch many-to-many field previous_locations on a sealed object." with self.assertRaisesMessage(SealedObject, message): instance.previous_locations.all() def test_sealed_instance_reverse_m2m_access(self): instance = Location.from_db('default', ['id'], [1]) - instance._state.sealed = True + instance.seal() message = "Cannot fetch many-to-many field previous_visitors on a sealed object." with self.assertRaisesMessage(SealedObject, message): instance.previous_visitors.all()