diff --git a/kobo/apps/openrosa/apps/api/tests/viewsets/test_xform_submission_api.py b/kobo/apps/openrosa/apps/api/tests/viewsets/test_xform_submission_api.py index cfbeabf0b0..662a704713 100644 --- a/kobo/apps/openrosa/apps/api/tests/viewsets/test_xform_submission_api.py +++ b/kobo/apps/openrosa/apps/api/tests/viewsets/test_xform_submission_api.py @@ -10,6 +10,7 @@ from constance.test import override_config from django.conf import settings from django.contrib.auth.models import AnonymousUser +from django.core.cache import cache from django.core.exceptions import PermissionDenied from django.core.files.base import ContentFile from django.core.files.uploadedfile import InMemoryUploadedFile @@ -144,7 +145,7 @@ def test_over_limit_submission_rejection_anonymous(self): def test_over_limit_submission_rejection_authenticated(self, mock_usage): """ Ensure submissions by an authenticated user are rejected if asset owner - is over their storage or submission limit and that check_exceeded_limit + is over their storage or submission limit and that check_exceeded_limits is run. """ path = os.path.join( @@ -167,15 +168,16 @@ def test_over_limit_submission_rejection_authenticated(self, mock_usage): } mock_usage.return_value = mock_balances with patch( - 'kobo.apps.openrosa.libs.utils.logger_tools.check_exceeded_limit', - return_value=None, + 'kobo.apps.openrosa.libs.utils.logger_tools.check_exceeded_limits', + return_value={}, ) as patched: request = self.factory.post('/submission', data, format='json') auth = DigestAuth('bob', 'bobbob') request.META.update(auth(request.META, response)) response = self.view(request, username=self.user.username) - patched.assert_any_call(self.user, UsageType.SUBMISSION) - patched.assert_any_call(self.user, UsageType.STORAGE_BYTES) + patched.assert_called_once_with( + self.user, [UsageType.SUBMISSION, UsageType.STORAGE_BYTES] + ) self.assertEqual(response.status_code, status.HTTP_402_PAYMENT_REQUIRED) mock_balances = { @@ -186,8 +188,8 @@ def test_over_limit_submission_rejection_authenticated(self, mock_usage): } mock_usage.return_value = mock_balances with patch( - 'kobo.apps.openrosa.libs.utils.logger_tools.check_exceeded_limit', - return_value=None, + 'kobo.apps.openrosa.libs.utils.logger_tools.check_exceeded_limits', + return_value={}, ) as patched: request = self.factory.post('/submission', data, format='json') response = self.view(request) @@ -196,10 +198,59 @@ def test_over_limit_submission_rejection_authenticated(self, mock_usage): auth = DigestAuth('bob', 'bobbob') request.META.update(auth(request.META, response)) response = self.view(request, username=self.user.username) - patched.assert_any_call(self.user, UsageType.SUBMISSION) - patched.assert_any_call(self.user, UsageType.STORAGE_BYTES) + patched.assert_called_once_with( + self.user, [UsageType.SUBMISSION, UsageType.STORAGE_BYTES] + ) self.assertEqual(response.status_code, status.HTTP_402_PAYMENT_REQUIRED) + @pytest.mark.skipif( + not settings.STRIPE_ENABLED, reason='Requires stripe functionality' + ) + @patch( + 'kobo.apps.openrosa.libs.utils.logger_tools.ServiceUsageCalculator.get_usage_balances' # noqa: E501 + ) + @patch('kobo.apps.stripe.utils.limit_enforcement._get_usage_balances') + def test_over_limit_rejection_counts_on_fresh_usage( + self, mock_fresh_usage, mock_cached_usage + ): + """ + Ensure a rejected submission records exceeded-limit counters from fresh + usage, not from the cached balances that triggered the rejection + """ + from kobo.apps.stripe.models import ExceededLimitCounter + + mock_cached_usage.return_value = { + UsageType.STORAGE_BYTES: {'exceeded': True}, + UsageType.SUBMISSION: None, + } + # e.g. the owner deleted attachments after the balances were cached + mock_fresh_usage.return_value = { + UsageType.STORAGE_BYTES: {'exceeded': False}, + UsageType.SUBMISSION: None, + } + for usage_type in (UsageType.SUBMISSION, UsageType.STORAGE_BYTES): + cache.delete(f'{self.user.id}_checked_exceeded_{usage_type}_limit') + path = os.path.join( + os.path.dirname(os.path.abspath(__file__)), + '..', + 'fixtures', + 'transport_submission.json', + ) + with open(path, 'rb') as f: + data = json.loads(f.read()) + + request = self.factory.post('/submission', data, format='json') + response = self.view(request) + self.assertEqual(response.status_code, 401) + request = self.factory.post('/submission', data, format='json') + auth = DigestAuth('bob', 'bobbob') + request.META.update(auth(request.META, response)) + response = self.view(request, username=self.user.username) + + self.assertEqual(response.status_code, status.HTTP_402_PAYMENT_REQUIRED) + mock_fresh_usage.assert_called_once() + assert not ExceededLimitCounter.objects.exists() + @pytest.mark.skipif( not settings.STRIPE_ENABLED, reason='Requires stripe functionality' ) diff --git a/kobo/apps/openrosa/apps/logger/tests/test_simple_submission.py b/kobo/apps/openrosa/apps/logger/tests/test_simple_submission.py index d5c8456469..5f2f2e3b6f 100644 --- a/kobo/apps/openrosa/apps/logger/tests/test_simple_submission.py +++ b/kobo/apps/openrosa/apps/logger/tests/test_simple_submission.py @@ -153,12 +153,13 @@ def test_corrupted_submission(self): ) def test_check_exceeded_limit_on_submission(self): with patch( - 'kobo.apps.openrosa.libs.utils.logger_tools.check_exceeded_limit', - return_value=None, + 'kobo.apps.openrosa.libs.utils.logger_tools.check_exceeded_limits', + return_value={}, ) as patched: self._submit_simple_yes() - patched.assert_any_call(self.user, UsageType.SUBMISSION) - patched.assert_any_call(self.user, UsageType.STORAGE_BYTES) + patched.assert_called_once_with( + self.user, [UsageType.SUBMISSION, UsageType.STORAGE_BYTES] + ) def test_parsed_instance_submitted_by_value(self): class MockRequest: diff --git a/kobo/apps/openrosa/libs/utils/logger_tools.py b/kobo/apps/openrosa/libs/utils/logger_tools.py index 867293a29f..26fb5df146 100644 --- a/kobo/apps/openrosa/libs/utils/logger_tools.py +++ b/kobo/apps/openrosa/libs/utils/logger_tools.py @@ -93,7 +93,7 @@ from kobo.apps.openrosa.libs.utils.model_tools import queryset_iterator, set_uuid from kobo.apps.openrosa.libs.utils.viewer_tools import get_mongo_userform_id from kobo.apps.organizations.constants import UsageType -from kobo.apps.stripe.utils.limit_enforcement import check_exceeded_limit +from kobo.apps.stripe.utils.limit_enforcement import check_exceeded_limits from kpi.constants import PERM_ADD_SUBMISSIONS, PERM_CHANGE_SUBMISSIONS from kpi.deployment_backends.kc_access.storage import ( default_kobocat_storage as default_storage, @@ -320,13 +320,17 @@ def create_instance( ): calculator = ServiceUsageCalculator(xform.user) balances = calculator.get_usage_balances() - for usage_type in [UsageType.STORAGE_BYTES, UsageType.SUBMISSION]: - balance = balances[usage_type] - if balance and balance['exceeded']: - check_exceeded_limit(xform.user, UsageType.SUBMISSION) - check_exceeded_limit(xform.user, UsageType.STORAGE_BYTES) + exceeded_types = [ + usage_type + for usage_type in [UsageType.STORAGE_BYTES, UsageType.SUBMISSION] + if balances[usage_type] and balances[usage_type]['exceeded'] + ] + if exceeded_types: + check_exceeded_limits( + xform.user, [UsageType.SUBMISSION, UsageType.STORAGE_BYTES] + ) - raise ExceededUsageLimitError({'type': usage_type}) + raise ExceededUsageLimitError({'type': exceeded_types[0]}) # get root uuid root_uuid, fallback_on_uuid = get_root_uuid_from_xml(xml) @@ -433,8 +437,10 @@ def create_instance( ) if settings.STRIPE_ENABLED: - check_exceeded_limit(xform.user, UsageType.SUBMISSION) - check_exceeded_limit(xform.user, UsageType.STORAGE_BYTES) + check_exceeded_limits( + xform.user, + [UsageType.SUBMISSION, UsageType.STORAGE_BYTES], + ) return instance diff --git a/kobo/apps/stripe/tests/test_stripe_utils.py b/kobo/apps/stripe/tests/test_stripe_utils.py index 9d0afa67f4..6c2851e280 100644 --- a/kobo/apps/stripe/tests/test_stripe_utils.py +++ b/kobo/apps/stripe/tests/test_stripe_utils.py @@ -37,6 +37,7 @@ ) from kobo.apps.stripe.utils.limit_enforcement import ( check_exceeded_limit, + check_exceeded_limits, update_or_remove_limit_counter, ) from kobo.apps.stripe.utils.manual_subscription import ( @@ -762,8 +763,8 @@ def test_check_exceeded_limit_adds_counters(self): # We want to test this function directly here, so we patch it out when # it is called on submission to avoid cache restrictions with patch( - 'kobo.apps.openrosa.libs.utils.logger_tools.check_exceeded_limit', - return_value=None, + 'kobo.apps.openrosa.libs.utils.logger_tools.check_exceeded_limits', + return_value={}, ): self.add_submissions(count=2, asset=self.asset, username='someuser') self.add_nlp_trackers() @@ -790,8 +791,8 @@ def test_check_exceeded_limit_updates_counters(self): # We want to test this function directly here, so we patch it out when # is called on submission to avoid cache restrictions with patch( - 'kobo.apps.openrosa.libs.utils.logger_tools.check_exceeded_limit', - return_value=None, + 'kobo.apps.openrosa.libs.utils.logger_tools.check_exceeded_limits', + return_value={}, ): self.add_submissions(count=2, asset=self.asset, username='someuser') self.add_nlp_trackers() @@ -866,3 +867,28 @@ def test_update_or_remove_limit_counter(self): ): update_or_remove_limit_counter(counter) assert ExceededLimitCounter.objects.count() == 0 + + def test_check_exceeded_limits_computes_balances_once(self): + mock_balances = { + UsageType.ASR_SECONDS: None, + UsageType.MT_CHARACTERS: None, + UsageType.STORAGE_BYTES: {'exceeded': True}, + UsageType.SUBMISSION: {'exceeded': True}, + } + with patch( + 'kpi.utils.usage_calculator.ServiceUsageCalculator.get_usage_balances', + return_value=mock_balances, + ) as patched: + check_exceeded_limits( + self.someuser, + [UsageType.SUBMISSION, UsageType.STORAGE_BYTES], + ) + patched.assert_called_once() + + for usage_type in [UsageType.SUBMISSION, UsageType.STORAGE_BYTES]: + assert ( + ExceededLimitCounter.objects.filter( + user_id=self.anotheruser.id, limit_type=usage_type + ).count() + == 1 + ) diff --git a/kobo/apps/stripe/utils/limit_enforcement.py b/kobo/apps/stripe/utils/limit_enforcement.py index dc67e079b6..ea9e3c9e1e 100644 --- a/kobo/apps/stripe/utils/limit_enforcement.py +++ b/kobo/apps/stripe/utils/limit_enforcement.py @@ -1,4 +1,5 @@ from datetime import timedelta +from typing import TYPE_CHECKING from django.conf import settings from django.core.cache import cache @@ -9,42 +10,71 @@ from kobo.apps.stripe.utils.import_management import requires_stripe from kpi.utils.usage_calculator import ServiceUsageCalculator +if TYPE_CHECKING: + from kobo.apps.stripe.models import ExceededLimitCounter + @requires_stripe def check_exceeded_limit(user, usage_type: UsageType, **kwargs): """ - Checks whether user's org has exceeded its limits for a given - usage type and updates ExceededLimitCounters accordingly. Uses - cached key to avoid running checks more than once within - the ENDPOINT_CACHE_DURATION + Check a single usage type against its limit. + + Thin wrapper around `check_exceeded_limits` kept for callers that work with + one usage type. Returns the created/updated counter, or None when the check + is gated or the limit is not exceeded. + """ + return check_exceeded_limits(user, [usage_type]).get(usage_type) + + +@requires_stripe +def check_exceeded_limits( + user, + usage_types: list[UsageType], + **kwargs, +) -> 'dict[UsageType, ExceededLimitCounter]': + """ + Check whether user's org has exceeded its limits for the given usage types + and update ExceededLimitCounters accordingly. + + Uses a cached key per usage type to avoid running checks more than once + within the ENDPOINT_CACHE_DURATION. Fresh usage balances are computed at most + once per call, and only when a usage type passes its gate. """ org = user.organization if org.is_mmo: user = org.owner_user_object - cache_key = f'{user.id}_checked_exceeded_{usage_type}_limit' + ExceededLimitCounter = kwargs['exceeded_limit_counter_model'] + balances = None + counters = {} + + for usage_type in usage_types: + cache_key = f'{user.id}_checked_exceeded_{usage_type}_limit' - if cache.get(cache_key, None): - return + if cache.get(cache_key, None): + continue - ExceededLimitCounter = kwargs['exceeded_limit_counter_model'] - balances = _get_usage_balances(user) - balance = balances[usage_type] - counter = None - - if balance and balance['exceeded']: - counter, created = ExceededLimitCounter.objects.get_or_create( - user=user, - limit_type=usage_type, - ) - - if not created and counter.date_modified.date() < timezone.now().date(): - delta = timezone.now().date() - counter.date_modified.date() - counter.days += delta.days - counter.save() - - cache.set(cache_key, True, settings.ENDPOINT_CACHE_DURATION) - return counter + if balances is None: + balances = _get_usage_balances(user) + + balance = balances[usage_type] + + if balance and balance['exceeded']: + counter, created = ExceededLimitCounter.objects.get_or_create( + user=user, + limit_type=usage_type, + ) + + if not created and counter.date_modified.date() < timezone.now().date(): + delta = timezone.now().date() - counter.date_modified.date() + counter.days += delta.days + counter.save() + + counters[usage_type] = counter + + cache.set(cache_key, True, settings.ENDPOINT_CACHE_DURATION) + + return counters @requires_stripe