From 164a6f5ded272694c86da07dde84fd4644afd3e8 Mon Sep 17 00:00:00 2001 From: uzairr Date: Mon, 18 Mar 2024 11:51:46 +0500 Subject: [PATCH] test: add tests for prefetch and select related qs --- course_discovery/apps/core/tests/test_utils.py | 6 ++++++ course_discovery/apps/core/utils.py | 4 ++-- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/course_discovery/apps/core/tests/test_utils.py b/course_discovery/apps/core/tests/test_utils.py index 08f71c3530..bc8596cd89 100644 --- 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 index afeb3fd5e5..9560f0de84 100644 --- 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