From 88e3459e02f004ab026434c106b8d1e7ca3b3da6 Mon Sep 17 00:00:00 2001 From: Muhammad Afaq Shuaib Date: Tue, 30 Apr 2024 03:58:24 +0500 Subject: [PATCH] test: add allow restricted_flag --- .../apps/course_metadata/managers.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/course_discovery/apps/course_metadata/managers.py b/course_discovery/apps/course_metadata/managers.py index 6e8e37432d..98de0f39f0 100644 --- a/course_discovery/apps/course_metadata/managers.py +++ b/course_discovery/apps/course_metadata/managers.py @@ -8,17 +8,23 @@ class DraftManager(models.Manager): """ Model manager that hides draft rows unless you ask for them. """ - def get_queryset(self): - return super().get_queryset().filter(draft=models.Value(0)) - - def _with_drafts(self): - return super().get_queryset() + def get_queryset(self, allow_restricted=False): + if allow_restricted: + return super().get_queryset().filter(draft=models.Value(0)) + return super().get_queryset().filter(draft=models.Value(0), restricted_run__isnull=True) + + def _with_drafts(self, allow_restricted=False): + if allow_restricted: + return super().get_queryset() + return super().get_queryset().filter(restricted_run__isnull=True) - def filter_drafts(self, **kwargs): + def filter_drafts(self, allow_restricted=False, **kwargs): """ Acts like filter(), but prefers draft versions. If a draft is not available, we give back the non-draft version. """ + if allow_restricted: + return self._with_drafts(allow_restricted=allow_restricted).filter(Q(draft=models.Value(1)) | Q(draft_version=None)).filter(**kwargs) return self._with_drafts().filter(Q(draft=models.Value(1)) | Q(draft_version=None)).filter(**kwargs) def get_draft(self, **kwargs):