diff --git a/futurex_openedx_extensions/helpers/apps.py b/futurex_openedx_extensions/helpers/apps.py index 9c084d05..d87b389a 100644 --- a/futurex_openedx_extensions/helpers/apps.py +++ b/futurex_openedx_extensions/helpers/apps.py @@ -29,4 +29,6 @@ class HelpersConfig(AppConfig): def ready(self) -> None: """Connect handlers to send notifications about discussions.""" + from futurex_openedx_extensions.helpers import \ + monkey_patches # pylint: disable=unused-import, import-outside-toplevel from futurex_openedx_extensions.helpers import signals # pylint: disable=unused-import, import-outside-toplevel diff --git a/futurex_openedx_extensions/helpers/monkey_patches.py b/futurex_openedx_extensions/helpers/monkey_patches.py new file mode 100644 index 00000000..d84436b5 --- /dev/null +++ b/futurex_openedx_extensions/helpers/monkey_patches.py @@ -0,0 +1,21 @@ +"""Monkey patches defined here.""" +from __future__ import annotations + +from typing import Any + +from django.db.models.query import QuerySet + +original_queryset_chain = QuerySet._chain # pylint: disable=protected-access + + +def customized_queryset_chain(self: Any, **kwargs: Any) -> QuerySet: + """Customized queryset chain method for the QuerySet class.""" + result = original_queryset_chain(self, **kwargs) + + if hasattr(self, 'removable_annotations'): + result.removable_annotations = self.removable_annotations.copy() + + return result + + +QuerySet._chain = customized_queryset_chain # pylint: disable=protected-access diff --git a/futurex_openedx_extensions/helpers/pagination.py b/futurex_openedx_extensions/helpers/pagination.py index 449c1330..61e6d341 100644 --- a/futurex_openedx_extensions/helpers/pagination.py +++ b/futurex_openedx_extensions/helpers/pagination.py @@ -1,9 +1,28 @@ """Pagination helpers and classes for the API views.""" +from django.core.paginator import Paginator +from django.db.models.query import QuerySet +from django.utils.functional import cached_property from rest_framework.pagination import PageNumberPagination +class DefaultPaginator(Paginator): + """Default paginator settings for the API views.""" + @cached_property + def count(self) -> int: + """Return the total number of objects, across all pages.""" + if isinstance(self.object_list, QuerySet) and hasattr(self.object_list, 'removable_annotations'): + clone = self.object_list._chain() # pylint: disable=protected-access + for key in self.object_list.removable_annotations: + clone.query.annotations.pop(key, None) + return clone.count() + + return super().count + + class DefaultPagination(PageNumberPagination): """Default pagination settings for the API views.""" page_size: int = 20 page_size_query_param: str = 'page_size' max_page_size: int = 100 + + django_paginator_class = DefaultPaginator diff --git a/futurex_openedx_extensions/helpers/querysets.py b/futurex_openedx_extensions/helpers/querysets.py index fed2d6fb..df12186c 100644 --- a/futurex_openedx_extensions/helpers/querysets.py +++ b/futurex_openedx_extensions/helpers/querysets.py @@ -18,6 +18,42 @@ from futurex_openedx_extensions.helpers.users import get_user_by_key +def update_removable_annotations( + queryset: QuerySet, + removable: set | List[str] | None = None, + not_removable: set | List[str] | None = None, +) -> None: + """ + Update the removable annotations on the given queryset. + + :param queryset: QuerySet to update + :type queryset: QuerySet + :param removable: Set of annotations to add to the removable annotations + :type removable: set(str) | List[str] | None + :param not_removable: Set of annotations to remove from the removable annotations + :type not_removable: set(str) | List[str] | None + """ + removable_annotations = queryset.removable_annotations if hasattr(queryset, 'removable_annotations') else set() + removable_annotations = (removable_annotations | set(removable or [])) - set(not_removable or []) + + if not removable_annotations and hasattr(queryset, 'removable_annotations'): + delattr(queryset, 'removable_annotations') + + elif removable_annotations: + queryset.removable_annotations = removable_annotations + + +def clear_removable_annotations(queryset: QuerySet) -> None: + """ + Clear the removable annotations on the given queryset. + + :param queryset: QuerySet to clear the removable annotations from + :type queryset: QuerySet + """ + if hasattr(queryset, 'removable_annotations'): + del queryset.removable_annotations + + def check_staff_exist_queryset( ref_user_id: str | Value, ref_org: str | Value | List | None, diff --git a/tests/test_helpers/test_pagination.py b/tests/test_helpers/test_pagination.py index 22f78c06..f6ab8e67 100644 --- a/tests/test_helpers/test_pagination.py +++ b/tests/test_helpers/test_pagination.py @@ -1,7 +1,10 @@ """Tests for pagination helpers""" +from unittest.mock import MagicMock, PropertyMock, patch + +from django.db.models import QuerySet from rest_framework.pagination import PageNumberPagination -from futurex_openedx_extensions.helpers.pagination import DefaultPagination +from futurex_openedx_extensions.helpers.pagination import DefaultPagination, DefaultPaginator def test_default_pagination(): @@ -10,3 +13,46 @@ def test_default_pagination(): assert DefaultPagination.page_size == 20 assert DefaultPagination.page_size_query_param == 'page_size' assert DefaultPagination.max_page_size == 100 + assert DefaultPagination.django_paginator_class == DefaultPaginator + + +@patch('futurex_openedx_extensions.helpers.pagination.Paginator.count', new_callable=PropertyMock) +def test_count_with_query_set_and_removable_annotations(mock_super_count): + mock_super_count.return_value = 'should not be reached' + + mock_queryset = MagicMock(spec=QuerySet) + mock_queryset.removable_annotations = {'annotation1'} + mock_queryset._chain.return_value = mock_queryset + mock_queryset.query.annotations = {'annotation1': None, 'annotation2': None} + mock_queryset.count.return_value = 5 + + paginator = DefaultPaginator(mock_queryset, per_page=10) + + assert paginator.count == 5 + assert 'annotation1' not in mock_queryset.query.annotations + mock_super_count.assert_not_called() + + +@patch('futurex_openedx_extensions.helpers.pagination.Paginator.count', new_callable=PropertyMock) +def test_count_with_query_set_no_removable_annotations(mock_super_count): + mock_super_count.return_value = 44 + + mock_queryset = MagicMock(spec=QuerySet) + delattr(mock_queryset, 'removable_annotations') + mock_queryset.count.return_value = 'should not be reached' + mock_queryset.query.annotations = {'annotation1': None, 'annotation2': None} + + paginator = DefaultPaginator(mock_queryset, per_page=10) + + assert paginator.count == mock_super_count.return_value + assert mock_queryset.query.annotations == {'annotation1': None, 'annotation2': None} + + +@patch('futurex_openedx_extensions.helpers.pagination.Paginator.count', new_callable=PropertyMock) +def test_count_with_not_query_set(mock_super_count): + mock_super_count.return_value = 44 + + object_list = 'not a queryset' + paginator = DefaultPaginator(object_list, per_page=10) + + assert paginator.count == mock_super_count.return_value diff --git a/tests/test_helpers/test_querysets.py b/tests/test_helpers/test_querysets.py index b9a6b7f0..5cd294a6 100644 --- a/tests/test_helpers/test_querysets.py +++ b/tests/test_helpers/test_querysets.py @@ -1,4 +1,6 @@ """Tests for querysets helpers""" +from unittest.mock import Mock + import pytest from common.djangoapps.student.models import CourseAccessRole, UserProfile from django.contrib.auth import get_user_model @@ -217,3 +219,40 @@ def test_get_permitted_learners_queryset( result = querysets.get_permitted_learners_queryset(queryset, fx_permission_info, include_staff=True) assert result.count() == expected_with_staff + + +@pytest.mark.parametrize('original, removable, not_removable, expected_result', [ + (None, None, None, None), + (None, {'f1', 'f2'}, None, {'f1', 'f2'}), + (None, None, {'f1', 'f2'}, None), + ({'f1', 'f2'}, None, None, {'f1', 'f2'}), + ({'f1', 'f2'}, {'f1'}, None, {'f1', 'f2'}), + ({'f1', 'f2'}, None, {'f1'}, {'f2'}), + ({'f1', 'f2'}, None, {'f1', 'f3'}, {'f2'}), + ({'f1', 'f2'}, {'f1'}, {'f1', 'f2'}, None), +]) +def test_update_removable_annotations(original, removable, not_removable, expected_result): + """Verify update_removable_annotations function.""" + queryset = Mock() + if original is None: + delattr(queryset, 'removable_annotations') + else: + queryset.removable_annotations = original + + querysets.update_removable_annotations(queryset, removable, not_removable) + + if expected_result is None: + assert not hasattr(queryset, 'removable_annotations') + else: + assert queryset.removable_annotations == expected_result + + +def test_clear_removable_annotations(): + """Verify clear_removable_annotations function.""" + queryset = Mock(removable_annotations={'f1', 'f2'}) + querysets.clear_removable_annotations(queryset) + + assert not hasattr(queryset, 'removable_annotations') + + querysets.clear_removable_annotations(queryset) + assert not hasattr(queryset, 'removable_annotations')