From 8c87f2f53bdf62043df233a1787bee33289cf13b Mon Sep 17 00:00:00 2001 From: Mike Date: Sun, 23 Jun 2024 08:52:45 -0700 Subject: [PATCH] fix: ensure default manager is correctly replaced and ordered --- eav/registry.py | 44 ++++++++++++++++++++++++++++------------- test_project/models.py | 45 ++++++++++++++++++++++++++++++++++++++++++ tests/test_registry.py | 21 ++++++++++++++++++++ 3 files changed, 96 insertions(+), 14 deletions(-) diff --git a/eav/registry.py b/eav/registry.py index 40879a6..8d2d1d7 100644 --- a/eav/registry.py +++ b/eav/registry.py @@ -102,25 +102,41 @@ class Registry(object): self.model_cls = model_cls self.config_cls = model_cls._eav_config_cls - def _attach_manager(self): + def _attach_manager(self) -> None: """ - Attach the manager to *manager_attr* specified in *config_cls* + Attach the EntityManager to the model class. + + This method replaces the existing manager specified in the `config_cls` + with a new instance of `EntityManager`. If the specified manager is the + default manager, the `EntityManager` is set as the new default manager. + Otherwise, it is appended to the list of managers. + + If the model class already has a manager with the same name as the one + specified in `config_cls`, it is saved as `old_mgr` in the `config_cls` + for use during detachment. """ - # Save the old manager if the attribute name conflicts with the new one. - if hasattr(self.model_cls, self.config_cls.manager_attr): - mgr = getattr(self.model_cls, self.config_cls.manager_attr) + manager_attr = self.config_cls.manager_attr + model_meta = self.model_cls._meta + current_manager = getattr(self.model_cls, manager_attr, None) - # For some models, `local_managers` may be empty, eg. - # django.contrib.auth.models.User and AbstractUser - if mgr in self.model_cls._meta.local_managers: - self.config_cls.old_mgr = mgr - self.model_cls._meta.local_managers.remove(mgr) + if isinstance(current_manager, EntityManager): + # EntityManager is already attached, no need to proceed + return - self.model_cls._meta._expire_cache() + # Create a new EntityManager + new_manager = EntityManager() - # Attach the new manager to the model. - mgr = EntityManager() - mgr.contribute_to_class(self.model_cls, self.config_cls.manager_attr) + # Save and remove the old manager if it exists + if current_manager and current_manager in model_meta.local_managers: + self.config_cls.old_mgr = current_manager + model_meta.local_managers.remove(current_manager) + + # Set the creation_counter to maintain the order + # This ensures that the new manager has the same priority as the old one + new_manager.creation_counter = current_manager.creation_counter + + # Attach the new EntityManager instance to the model. + new_manager.contribute_to_class(self.model_cls, manager_attr) def _detach_manager(self): """ diff --git a/test_project/models.py b/test_project/models.py index 4e8b33b..00e0078 100644 --- a/test_project/models.py +++ b/test_project/models.py @@ -10,6 +10,7 @@ from django.db import models from eav.decorators import register_eav from eav.models import EAVModelMeta +from eav.managers import EntityManager #: Constants MAX_CHARFIELD_LEN: Final = 254 @@ -25,6 +26,47 @@ class TestBase(models.Model): abstract = True +class DoctorManager(EntityManager): + """ + Custom manager for the Doctor model. + + This manager extends the EntityManager and provides additional + methods specific to the Doctor model, and is expected to be the + default manager on the model. + """ + + def get_by_name(self, name: str) -> models.QuerySet: + """Returns a QuerySet of doctors with the given name. + + Args: + name (str): The name of the doctor to search for. + + Returns: + models.QuerySet: A QuerySet of doctors with the specified name. + """ + return self.filter(name=name) + + +class DoctorSubstringManager(models.Manager): + """ + Custom manager for the Doctor model. + + This is a second manager used to ensure during testing that it's not replaced + as the default manager after eav.register(). + """ + + def get_by_name_contains(self, substring: str) -> models.QuerySet: + """Returns a QuerySet of doctors whose names contain the given substring. + + Args: + substring (str): The substring to search for in the doctor's name. + + Returns: + models.QuerySet: A QuerySet of doctors whose names contain the specified substring. + """ + return self.filter(name__icontains=substring) + + @final @register_eav() class Doctor(TestBase): @@ -33,6 +75,9 @@ class Doctor(TestBase): id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False) name = models.CharField(max_length=MAX_CHARFIELD_LEN) + objects = DoctorManager() + substrings = DoctorSubstringManager() + @final class Patient(TestBase): diff --git a/tests/test_registry.py b/tests/test_registry.py index 877e568..663cd4b 100644 --- a/tests/test_registry.py +++ b/tests/test_registry.py @@ -2,8 +2,10 @@ from django.contrib.auth.models import User from django.test import TestCase import eav +from eav.managers import EntityManager from eav.registry import EavConfig from test_project.models import ( + Doctor, Encounter, ExampleMetaclassModel, ExampleModel, @@ -112,3 +114,22 @@ class RegistryTests(TestCase): # Reverse check: managers should be empty again eav.unregister(User) assert bool(User._meta.local_managers) is False + +def test_default_manager_stays() -> None: + """ + Test to ensure default manager remains after registration. + + This test verifies that the default manager of the Doctor model is correctly + replaced or maintained after registering a new EntityManager. Specifically, + if the model's Meta default_manager_name isn't set, the test ensures that + the default manager remains as 'objects' or the first manager declared in + the class. + """ + instance_meta = Doctor._meta + assert instance_meta.default_manager_name is None + assert isinstance(instance_meta.default_manager, EntityManager) + + # Explicity test this as for our test setup, we want to have a state where + # the default manager is 'objects' + assert instance_meta.default_manager.name == 'objects' + assert len(instance_meta.managers) == 2