diff --git a/generic_links/managers.py b/generic_links/managers.py new file mode 100644 index 0000000..5771ea4 --- /dev/null +++ b/generic_links/managers.py @@ -0,0 +1,12 @@ +from django.db import models + + +class GenericLinkQuerySet(models.QuerySet): + def internal(self): + return self.exclude(is_external=True) + + def external(self): + return self.filter(is_external=True) + + def by_user(self, user): + return self.filter(user=user) diff --git a/generic_links/models.py b/generic_links/models.py index 05d6640..fef77ed 100644 --- a/generic_links/models.py +++ b/generic_links/models.py @@ -6,6 +6,8 @@ from django.db import models from django.utils.translation import gettext_lazy as _ +from .managers import GenericLinkQuerySet + class GenericLink(models.Model): """ @@ -25,6 +27,8 @@ class GenericLink(models.Model): is_external = models.BooleanField(default=True, db_index=True) + objects = GenericLinkQuerySet.as_manager() + class Meta: ordering = ("-created_at",) verbose_name = _("Generic Link") diff --git a/tests/test_managers.py b/tests/test_managers.py new file mode 100644 index 0000000..dc682ac --- /dev/null +++ b/tests/test_managers.py @@ -0,0 +1,65 @@ +from django.contrib.auth.models import User +from django.test import TestCase + +from generic_links.models import GenericLink + + +class ModelTest(TestCase): + def test_external_and_internal(self): + user = User.objects.create(username="Test User") + + external_link = GenericLink.objects.create( + title="An external link", + url="https://www.test.com", + description="An external link for this user", + content_object=user, + is_external=True, + ) + + internal_link = GenericLink.objects.create( + title="An internal link", + url="https://www.test.com", + description="An internal link for this user", + content_object=user, + is_external=False, + ) + + self.assertEqual(GenericLink.objects.count(), 2) + + external_qs = GenericLink.objects.external() + + self.assertEqual(external_qs.count(), 1) + self.assertListEqual([link.id for link in external_qs], [external_link.id]) + + internal_qs = GenericLink.objects.internal() + + self.assertEqual(internal_qs.count(), 1) + self.assertListEqual([link.id for link in internal_qs], [internal_link.id]) + + def test_by_user(self): + user1 = User.objects.create(username="user1") + + link = GenericLink.objects.create( + title="Test Title", + url="https://www.test.com", + description="Test Description", + content_object=user1, + user=user1, + ) + + user2 = User.objects.create(username="user2") + + link2 = GenericLink.objects.create( + title="Second Test Title", + url="https://www.test.com", + description="Test Description for user 2", + content_object=user2, + user=user1, + ) + + self.assertEqual(GenericLink.objects.count(), 2) + + self.assertEqual(GenericLink.objects.by_user(user1).count(), 2) + self.assertSetEqual({link.id for link in GenericLink.objects.by_user(user1)}, {link.id, link2.id}) + + self.assertEqual(GenericLink.objects.by_user(user2).count(), 0)