Custom managers for filtering by owner, and internal vs external links

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)

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")

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)