test: add tests for prefetch and select related qs

iamsobanjaved · Apr 1, 2024 · 164a6f5 · 164a6f5
1 parent 81842b3
commit 164a6f5
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/course_discovery/apps/core/tests/test_utils.py b/course_discovery/apps/core/tests/test_utils.py
@@ -121,3 +121,9 @@ def test_iter(self):
 
     def test_getitem(self):
         assert self.course_runs[0] == self.wrapper[0]
+
+    def test_prefetch_related_count(self):
+        assert self.search_queryset.count() == self.wrapper.prefetch_related().count()
+
+    def test_select_related_count(self):
+        assert self.search_queryset.count() == self.wrapper.select_related().count()
diff --git a/course_discovery/apps/core/utils.py b/course_discovery/apps/core/utils.py
@@ -164,7 +164,7 @@ def __init__(self, queryset, model):
     def prefetch_related(self, *lookups):
         """Same as QuerySet.prefetch_related()"""
         clone = self._chain()
-        if lookups == (None,):
+        if not lookups or lookups == (None,):
             clone._prefetch_related_lookups = ()  # pylint: disable=protected-access
         else:
             clone._prefetch_related_lookups += lookups
@@ -173,7 +173,7 @@ def prefetch_related(self, *lookups):
     def select_related(self, *lookups):
         """Will work same as .prefetch_related()"""
         clone = self._chain()
-        if lookups == (None,):
+        if not lookups or lookups == (None,):
             clone._select_related_lookups = ()  # pylint: disable=protected-access
         else:
             clone._select_related_lookups += lookups