From fa2fd238f86bf2b691dd8c6bdb380647cde9c9ed Mon Sep 17 00:00:00 2001 From: rgraber Date: Wed, 21 Aug 2024 16:12:57 -0400 Subject: [PATCH 01/10] fix: fix failing unit tests --- .../audit_log/tests/test_one_time_auth.py | 27 ++++++++----------- kpi/authentication.py | 2 +- 2 files changed, 12 insertions(+), 17 deletions(-) diff --git a/kobo/apps/audit_log/tests/test_one_time_auth.py b/kobo/apps/audit_log/tests/test_one_time_auth.py index 172f269517..383e870833 100644 --- a/kobo/apps/audit_log/tests/test_one_time_auth.py +++ b/kobo/apps/audit_log/tests/test_one_time_auth.py @@ -47,39 +47,34 @@ def setUp(self): # expected authentication type, method that needs to be mocked, endpoint to hit # (kpi and openrosa endpoints use different auth methods, and we want to test endpoints in both v1 and v2) ( - 'Token', + 'token', 'kpi.authentication.DRFTokenAuthentication.authenticate', 'data-list', ), ( - 'Basic', + 'basic', 'kpi.authentication.DRFBasicAuthentication.authenticate', 'api_v2:audit-log-list', ), ( - 'OAuth2', + 'oauth2', 'kpi.authentication.OPOAuth2Authentication.authenticate', 'data-list', ), ( - 'Https Basic', + 'https basic', 'kobo.apps.openrosa.libs.authentication.BasicAuthentication.authenticate', 'data-list', ), ( - 'Token', + 'token', 'kpi.authentication.DRFTokenAuthentication.authenticate', - 'api_v2:submission-list', + 'api_v2:asset-list', ), ( - 'OAuth2', + 'oauth2', 'kpi.authentication.OPOAuth2Authentication.authenticate', - 'api_v2:submission-list', - ), - ( - 'Https Basic', - 'kobo.apps.openrosa.libs.authentication.BasicAuthentication.authenticate', - 'api_v2:submission-list', + 'api_v2:asset-list', ), ) @unpack @@ -121,11 +116,11 @@ def side_effect(request): return_value=True, side_effect=side_effect, ): - self.client.get(reverse('api_v2:submission-list'), **header) + self.client.get(reverse('api_v2:asset-list'), **header) log_exists = AuditLog.objects.filter( user_uid=TestOneTimeAuthentication.user.extra_details.uid, action=AuditAction.AUTH, - metadata__auth_type='Digest', + metadata__auth_type='digest', ).exists() self.assertTrue(log_exists) self.assertEqual(AuditLog.objects.count(), 1) @@ -154,7 +149,7 @@ def side_effect(request): log_exists = AuditLog.objects.filter( user_uid=TestOneTimeAuthentication.user.extra_details.uid, action=AuditAction.AUTH, - metadata__auth_type='Submission', + metadata__auth_type='submission', ).exists() self.assertTrue(log_exists) self.assertEqual(AuditLog.objects.count(), 1) diff --git a/kpi/authentication.py b/kpi/authentication.py index dfb72003d7..8f358b0566 100644 --- a/kpi/authentication.py +++ b/kpi/authentication.py @@ -170,5 +170,5 @@ def authenticate(self, request): if result is None: return result user, creds = result - self.create_access_log(request, user, 'OAuth2') + self.create_access_log(request, user, 'oauth2') return user, creds From 55451698c53691f07357ebe46fd16f09f8e19bd2 Mon Sep 17 00:00:00 2001 From: rgraber Date: Wed, 21 Aug 2024 16:20:29 -0400 Subject: [PATCH 02/10] fixup!: fix failing unit tests --- kobo/apps/audit_log/tests/test_one_time_auth.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/kobo/apps/audit_log/tests/test_one_time_auth.py b/kobo/apps/audit_log/tests/test_one_time_auth.py index 383e870833..d3db488dea 100644 --- a/kobo/apps/audit_log/tests/test_one_time_auth.py +++ b/kobo/apps/audit_log/tests/test_one_time_auth.py @@ -116,7 +116,7 @@ def side_effect(request): return_value=True, side_effect=side_effect, ): - self.client.get(reverse('api_v2:asset-list'), **header) + self.client.get(reverse('data-list'), **header) log_exists = AuditLog.objects.filter( user_uid=TestOneTimeAuthentication.user.extra_details.uid, action=AuditAction.AUTH, From abcda2be1faa188f3892635d8ae6eb3f4c1d62c5 Mon Sep 17 00:00:00 2001 From: Guillermo Date: Wed, 21 Aug 2024 21:23:27 -0600 Subject: [PATCH 03/10] Refactor ServiceUsageSerializer into a reusable utility to calculate usage numbers --- .../stripe/tests/test_organization_usage.py | 6 +- kpi/serializers/v2/service_usage.py | 164 ++--------- kpi/tests/api/v2/test_api_service_usage.py | 193 +------------ kpi/tests/test_usage_calculator.py | 267 ++++++++++++++++++ kpi/utils/usage_calculator.py | 146 ++++++++++ 5 files changed, 439 insertions(+), 337 deletions(-) create mode 100644 kpi/tests/test_usage_calculator.py create mode 100644 kpi/utils/usage_calculator.py diff --git a/kobo/apps/stripe/tests/test_organization_usage.py b/kobo/apps/stripe/tests/test_organization_usage.py index 0d234ce8f2..9a5af8a1ea 100644 --- a/kobo/apps/stripe/tests/test_organization_usage.py +++ b/kobo/apps/stripe/tests/test_organization_usage.py @@ -15,13 +15,13 @@ from kobo.apps.organizations.models import Organization, OrganizationUser from kobo.apps.stripe.tests.utils import generate_enterprise_subscription, generate_plan_subscription from kobo.apps.trackers.tests.submission_utils import create_mock_assets, add_mock_submissions -from kpi.tests.api.v2.test_api_service_usage import ServiceUsageAPIBase +from kpi.tests.test_usage_calculator import BaseUsageCalculatorTestCase from kpi.tests.api.v2.test_api_asset_usage import AssetUsageAPITestCase from rest_framework import status -class OrganizationServiceUsageAPIMultiUserTestCase(ServiceUsageAPIBase): +class OrganizationServiceUsageAPIMultiUserTestCase(BaseUsageCalculatorTestCase): """ Test organization service usage when Stripe is enabled. @@ -149,7 +149,7 @@ def test_endpoint_is_cached(self): self.expected_file_size() * self.expected_submissions_multi ) -class OrganizationServiceUsageAPITestCase(ServiceUsageAPIBase): +class OrganizationServiceUsageAPITestCase(BaseUsageCalculatorTestCase): org_id = 'orgAKWMFskafsngf' @classmethod diff --git a/kpi/serializers/v2/service_usage.py b/kpi/serializers/v2/service_usage.py index a3f65971b3..822f8ffe51 100644 --- a/kpi/serializers/v2/service_usage.py +++ b/kpi/serializers/v2/service_usage.py @@ -7,7 +7,10 @@ from kobo.apps.kobo_auth.shortcuts import User from kobo.apps.organizations.models import Organization -from kobo.apps.organizations.utils import get_monthly_billing_dates, get_yearly_billing_dates +from kobo.apps.organizations.utils import ( + get_monthly_billing_dates, + get_yearly_billing_dates, +) from kobo.apps.stripe.constants import ACTIVE_STRIPE_STATUSES from kobo.apps.trackers.models import NLPUsageCounter from kpi.deployment_backends.kc_access.shadow_models import ( @@ -16,6 +19,7 @@ ) from kpi.deployment_backends.kobocat_backend import KobocatDeploymentBackend from kpi.models.asset import Asset +from kpi.utils.usage_calculator import UsageCalculator class AssetUsageSerializer(serializers.HyperlinkedModelSerializer): @@ -117,158 +121,32 @@ class ServiceUsageSerializer(serializers.Serializer): def __init__(self, instance=None, data=empty, **kwargs): super().__init__(instance=instance, data=data, **kwargs) - - self._total_nlp_usage = {} - self._total_storage_bytes = 0 - self._total_submission_count = {} - self._current_month_start = None - self._current_month_end = None - self._current_year_start = None - self._current_year_end = None - self._organization = None - self._now = timezone.now() - self._get_per_asset_usage(instance) + organization = None + organization_id = self.context.get('organization_id', None) + if organization_id: + organization = Organization.objects.filter( + organization_users__user_id=instance.id, + id=organization_id, + ).first() + self.calculator = UsageCalculator(instance, organization) def get_total_nlp_usage(self, user): - return self._total_nlp_usage + return self.calculator.get_nlp_usage_counters() def get_total_submission_count(self, user): - return self._total_submission_count + return self.calculator.get_submission_counters() def get_total_storage_bytes(self, user): - return self._total_storage_bytes + return self.calculator.get_storage_usage() def get_current_month_start(self, user): - return self._current_month_start.isoformat() - + return self.calculator.current_month_start.isoformat() + def get_current_month_end(self, user): - return self._current_month_end.isoformat() + return self.calculator.current_month_end.isoformat() def get_current_year_start(self, user): - return self._current_year_start.isoformat() + return self.calculator.current_year_start.isoformat() def get_current_year_end(self, user): - return self._current_year_end.isoformat() - - def _filter_by_user(self, user_ids: list) -> Q: - """ - Turns a list of user ids into a query object to filter by - """ - return Q(user_id__in=user_ids) - - def _get_nlp_user_counters(self, month_filter, year_filter): - nlp_tracking = NLPUsageCounter.objects.only( - 'date', 'total_asr_seconds', 'total_mt_characters' - ).filter(self._user_id_query).aggregate( - asr_seconds_current_year=Coalesce( - Sum('total_asr_seconds', filter=year_filter), 0 - ), - mt_characters_current_year=Coalesce( - Sum('total_mt_characters', filter=year_filter), 0 - ), - asr_seconds_current_month=Coalesce( - Sum('total_asr_seconds', filter=month_filter), 0 - ), - mt_characters_current_month=Coalesce( - Sum('total_mt_characters', filter=month_filter), 0 - ), - asr_seconds_all_time=Coalesce(Sum('total_asr_seconds'), 0), - mt_characters_all_time=Coalesce(Sum('total_mt_characters'), 0), - ) - - for nlp_key, count in nlp_tracking.items(): - self._total_nlp_usage[nlp_key] = count if count is not None else 0 - - def _get_organization_details(self, user_id: int): - # Get the organization ID from the request - organization_id = self.context.get( - 'organization_id', None - ) - - if not organization_id: - return - - self._organization = Organization.objects.filter( - organization_users__user_id=user_id, - id=organization_id, - ).first() - - if not self._organization: - # Couldn't find organization, proceed as normal - return - - if settings.STRIPE_ENABLED: - # if the user is in an organization and has an enterprise plan, get all org users - # we evaluate this queryset instead of using it as a subquery because it's referencing - # fields from the auth_user tables on kpi *and* kobocat, making getting results in a - # single query not feasible until those tables are combined - user_ids = list( - User.objects.filter( - organizations_organization__id=organization_id, - organizations_organization__djstripe_customers__subscriptions__status__in=ACTIVE_STRIPE_STATUSES, - organizations_organization__djstripe_customers__subscriptions__items__price__product__metadata__has_key='plan_type', - organizations_organization__djstripe_customers__subscriptions__items__price__product__metadata__plan_type='enterprise', - ).values_list('pk', flat=True)[:settings.ORGANIZATION_USER_LIMIT] - ) - if user_ids: - self._user_id_query = self._filter_by_user(user_ids) - - def _get_per_asset_usage(self, user): - self._user_id = user.pk - self._user_id_query = self._filter_by_user([self._user_id]) - # get the billing data and list of organization users (if applicable) - self._get_organization_details(self._user_id) - - self._get_storage_usage() - - self._current_month_start, self._current_month_end = get_monthly_billing_dates(self._organization) - self._current_year_start, self._current_year_end = get_yearly_billing_dates(self._organization) - - current_month_filter = Q( - date__range=[self._current_month_start, self._now] - ) - current_year_filter = Q( - date__range=[self._current_year_start, self._now] - ) - - self._get_submission_counters(current_month_filter, current_year_filter) - self._get_nlp_user_counters(current_month_filter, current_year_filter) - - def _get_storage_usage(self): - """ - Get the storage used by non-(soft-)deleted projects for all users - - Users are represented by their ids with `self._user_ids` - """ - xforms = KobocatXForm.objects.only('attachment_storage_bytes', 'id').exclude( - pending_delete=True - ).filter(self._user_id_query) - - total_storage_bytes = xforms.aggregate( - bytes_sum=Coalesce(Sum('attachment_storage_bytes'), 0), - ) - - self._total_storage_bytes = total_storage_bytes['bytes_sum'] or 0 - - def _get_submission_counters(self, month_filter, year_filter): - """ - Calculate submissions for all users' projects even their deleted ones - - Users are represented by their ids with `self._user_ids` - """ - submission_count = KobocatDailyXFormSubmissionCounter.objects.only( - 'counter', 'user_id' - ).filter(self._user_id_query).aggregate( - all_time=Coalesce(Sum('counter'), 0), - current_year=Coalesce( - Sum('counter', filter=year_filter), 0 - ), - current_month=Coalesce( - Sum('counter', filter=month_filter), 0 - ), - ) - - for submission_key, count in submission_count.items(): - self._total_submission_count[submission_key] = ( - count if count is not None else 0 - ) + return self.calculator.current_year_end.isoformat() diff --git a/kpi/tests/api/v2/test_api_service_usage.py b/kpi/tests/api/v2/test_api_service_usage.py index 63826cd516..af5deb90f8 100644 --- a/kpi/tests/api/v2/test_api_service_usage.py +++ b/kpi/tests/api/v2/test_api_service_usage.py @@ -15,199 +15,10 @@ ) from kobo.apps.trackers.models import NLPUsageCounter from kpi.models import Asset -from kpi.tests.base_test_case import BaseAssetTestCase -from kpi.urls.router_api_v2 import URL_NAMESPACE as ROUTER_URL_NAMESPACE +from kpi.tests.test_usage_calculator import BaseUsageCalculatorTestCase -class ServiceUsageAPIBase(BaseAssetTestCase): - """ - This class contains setup logic and utility functions to test submissions/usage - """ - fixtures = ['test_data'] - - URL_NAMESPACE = ROUTER_URL_NAMESPACE - - xform = None - counter = None - attachment_id = 0 - - @classmethod - def setUpTestData(cls): - super().setUpTestData() - cls.anotheruser = User.objects.get(username='anotheruser') - cls.someuser = User.objects.get(username='someuser') - - def setUp(self): - super().setUp() - self.client.login(username='anotheruser', password='anotheruser') - - def _create_asset(self, user=None): - owner = user or self.anotheruser - content_source_asset = { - 'survey': [ - { - 'type': 'audio', - 'label': 'q1', - 'required': 'false', - '$kuid': 'abcd', - }, - { - 'type': 'file', - 'label': 'q2', - 'required': 'false', - '$kuid': 'efgh', - }, - ] - } - self.asset = Asset.objects.create( - content=content_source_asset, - owner=owner, - asset_type='survey', - ) - - self.asset.deploy(backend='mock', active=True) - self.asset.save() - - self.asset.deployment.set_namespace(self.URL_NAMESPACE) - self.submission_list_url = self.asset.deployment.submission_list_url - self._deployment = self.asset.deployment - - def add_nlp_trackers(self): - """ - Add nlp data to an asset - """ - # this month - today = timezone.now().date() - counter_1 = { - 'google_asr_seconds': 4586, - 'google_mt_characters': 5473, - } - NLPUsageCounter.objects.create( - user_id=self.anotheruser.id, - asset_id=self.asset.id, - date=today, - counters=counter_1, - total_asr_seconds=counter_1['google_asr_seconds'], - total_mt_characters=counter_1['google_mt_characters'], - ) - - # last month - last_month = today - relativedelta(months=1) - counter_2 = { - 'google_asr_seconds': 142, - 'google_mt_characters': 1253, - } - NLPUsageCounter.objects.create( - user_id=self.anotheruser.id, - asset_id=self.asset.id, - date=last_month, - counters=counter_2, - total_asr_seconds=counter_2['google_asr_seconds'], - total_mt_characters=counter_2['google_mt_characters'], - ) - - def add_submissions(self, count=2): - """ - Add one or more submissions to an asset (TWO by default) - """ - submissions = [] - v_uid = self.asset.latest_deployed_version.uid - - for x in range(count): - submission = { - '__version__': v_uid, - 'q1': 'audio_conversion_test_clip.3gp', - 'q2': 'audio_conversion_test_image.jpg', - '_uuid': str(uuid.uuid4()), - '_attachments': [ - { - 'id': self.attachment_id, - 'download_url': 'http://testserver/anotheruser/audio_conversion_test_clip.3gp', - 'filename': 'anotheruser/audio_conversion_test_clip.3gp', - 'mimetype': 'video/3gpp', - }, - { - 'id': self.attachment_id + 1, - 'download_url': 'http://testserver/anotheruser/audio_conversion_test_image.jpg', - 'filename': 'anotheruser/audio_conversion_test_image.jpg', - 'mimetype': 'image/jpeg', - }, - ], - '_submitted_by': 'anotheruser', - } - # increment the attachment ID for each attachment created - self.attachment_id = self.attachment_id + 2 - submissions.append(submission) - - self.asset.deployment.mock_submissions(submissions, flush_db=False) - self.update_xform_counters(self.asset, submissions=count) - - def update_xform_counters(self, asset: Asset, submissions: int = 0): - """ - Create/update the daily submission counter and the shadow xform we use to query it - """ - today = timezone.now() - if self.xform: - self.xform.attachment_storage_bytes += ( - self.expected_file_size() * submissions - ) - self.xform.save() - else: - xform_xml = ( - f'' - f'' - f'' - f' XForm test' - f' ' - f' ' - f' <{asset.uid} id="{asset.uid}" />' - f' ' - f' ' - f'' - f'' - f'' - f'' - ) - - self.xform = XForm.objects.create( - attachment_storage_bytes=( - self.expected_file_size() * submissions - ), - kpi_asset_uid=asset.uid, - date_created=today, - date_modified=today, - user_id=asset.owner_id, - xml=xform_xml, - json={} - ) - self.xform.save() - - if self.counter: - self.counter.counter += submissions - self.counter.save() - else: - self.counter = ( - DailyXFormSubmissionCounter.objects.create( - date=today.date(), - counter=submissions, - xform=self.xform, - user_id=asset.owner_id, - ) - ) - self.counter.save() - - def expected_file_size(self): - """ - Calculate the expected combined file size for the test audio clip and image - """ - return os.path.getsize( - settings.BASE_DIR + '/kpi/tests/audio_conversion_test_clip.3gp' - ) + os.path.getsize( - settings.BASE_DIR + '/kpi/tests/audio_conversion_test_image.jpg' - ) - - -class ServiceUsageAPITestCase(ServiceUsageAPIBase): +class ServiceUsageAPITestCase(BaseUsageCalculatorTestCase): def test_anonymous_user(self): """ Test that the endpoint is forbidden to anonymous user diff --git a/kpi/tests/test_usage_calculator.py b/kpi/tests/test_usage_calculator.py new file mode 100644 index 0000000000..0cc6236e8e --- /dev/null +++ b/kpi/tests/test_usage_calculator.py @@ -0,0 +1,267 @@ +import os.path +import uuid + +from dateutil.relativedelta import relativedelta +from django.conf import settings +from django.test import override_settings +from django.utils import timezone +from model_bakery import baker + +from kobo.apps.kobo_auth.shortcuts import User +from kobo.apps.openrosa.apps.logger.models import ( + XForm, + DailyXFormSubmissionCounter, +) +from kobo.apps.organizations.models import Organization +from kobo.apps.stripe.tests.utils import generate_enterprise_subscription +from kobo.apps.trackers.models import NLPUsageCounter +from kpi.models import Asset +from kpi.tests.base_test_case import BaseAssetTestCase +from kpi.utils.usage_calculator import UsageCalculator +from kpi.urls.router_api_v2 import URL_NAMESPACE as ROUTER_URL_NAMESPACE + + +class BaseUsageCalculatorTestCase(BaseAssetTestCase): + """ + This class contains setup logic and utility functions to test usage + calculations + """ + + fixtures = ['test_data'] + attachment_id = 0 + xform = None + counter = None + + URL_NAMESPACE = ROUTER_URL_NAMESPACE + + def setUp(self): + super().setUp() + self.client.login(username='anotheruser', password='anotheruser') + + @classmethod + def setUpTestData(cls): + super().setUpTestData() + cls.anotheruser = User.objects.get(username='anotheruser') + cls.someuser = User.objects.get(username='someuser') + + def _create_asset(self, user=None): + owner = user or self.anotheruser + content_source_asset = { + 'survey': [ + { + 'type': 'audio', + 'label': 'q1', + 'required': 'false', + '$kuid': 'abcd', + }, + { + 'type': 'file', + 'label': 'q2', + 'required': 'false', + '$kuid': 'efgh', + }, + ] + } + self.asset = Asset.objects.create( + content=content_source_asset, + owner=owner, + asset_type='survey', + ) + + self.asset.deploy(backend='mock', active=True) + self.asset.save() + + self.asset.deployment.set_namespace(self.URL_NAMESPACE) + self.submission_list_url = self.asset.deployment.submission_list_url + self._deployment = self.asset.deployment + + def add_nlp_trackers(self): + """ + Add nlp data to an asset + """ + # this month + today = timezone.now().date() + counter_1 = { + 'google_asr_seconds': 4586, + 'google_mt_characters': 5473, + } + NLPUsageCounter.objects.create( + user_id=self.anotheruser.id, + asset_id=self.asset.id, + date=today, + counters=counter_1, + total_asr_seconds=counter_1['google_asr_seconds'], + total_mt_characters=counter_1['google_mt_characters'], + ) + + # last month + last_month = today - relativedelta(months=1) + counter_2 = { + 'google_asr_seconds': 142, + 'google_mt_characters': 1253, + } + NLPUsageCounter.objects.create( + user_id=self.anotheruser.id, + asset_id=self.asset.id, + date=last_month, + counters=counter_2, + total_asr_seconds=counter_2['google_asr_seconds'], + total_mt_characters=counter_2['google_mt_characters'], + ) + + def add_submissions(self, count=2): + """ + Add one or more submissions to an asset (TWO by default) + """ + submissions = [] + v_uid = self.asset.latest_deployed_version.uid + + for x in range(count): + submission = { + '__version__': v_uid, + 'q1': 'audio_conversion_test_clip.3gp', + 'q2': 'audio_conversion_test_image.jpg', + '_uuid': str(uuid.uuid4()), + '_attachments': [ + { + 'id': self.attachment_id, + 'download_url': 'http://testserver/anotheruser/audio_conversion_test_clip.3gp', + 'filename': 'anotheruser/audio_conversion_test_clip.3gp', + 'mimetype': 'video/3gpp', + }, + { + 'id': self.attachment_id + 1, + 'download_url': 'http://testserver/anotheruser/audio_conversion_test_image.jpg', + 'filename': 'anotheruser/audio_conversion_test_image.jpg', + 'mimetype': 'image/jpeg', + }, + ], + '_submitted_by': 'anotheruser', + } + # increment the attachment ID for each attachment created + self.attachment_id = self.attachment_id + 2 + submissions.append(submission) + + self.asset.deployment.mock_submissions(submissions, flush_db=False) + self.update_xform_counters(self.asset, submissions=count) + + def expected_file_size(self): + """ + Calculate the expected combined file size for the test audio clip and image + """ + return os.path.getsize( + settings.BASE_DIR + '/kpi/tests/audio_conversion_test_clip.3gp' + ) + os.path.getsize( + settings.BASE_DIR + '/kpi/tests/audio_conversion_test_image.jpg' + ) + + def update_xform_counters(self, asset: Asset, submissions: int = 0): + """ + Create/update the daily submission counter and the shadow xform we use to query it + """ + today = timezone.now() + if self.xform: + self.xform.attachment_storage_bytes += ( + self.expected_file_size() * submissions + ) + self.xform.save() + else: + xform_xml = ( + f'' + f'' + f'' + f' XForm test' + f' ' + f' ' + f' <{asset.uid} id="{asset.uid}" />' + f' ' + f' ' + f'' + f'' + f'' + f'' + ) + + self.xform = XForm.objects.create( + attachment_storage_bytes=( + self.expected_file_size() * submissions + ), + kpi_asset_uid=asset.uid, + date_created=today, + date_modified=today, + user_id=asset.owner_id, + xml=xform_xml, + json={}, + ) + self.xform.save() + + if self.counter: + self.counter.counter += submissions + self.counter.save() + else: + self.counter = DailyXFormSubmissionCounter.objects.create( + date=today.date(), + counter=submissions, + xform=self.xform, + user_id=asset.owner_id, + ) + self.counter.save() + + +class UsageCalculatorTestCase(BaseUsageCalculatorTestCase): + def setUp(self): + super().setUp() + self._create_asset() + self.add_nlp_trackers() + self.add_submissions(count=5) + + def test_nlp_usage_counters(self): + calculator = UsageCalculator(self.anotheruser, None) + nlp_usage = calculator.get_nlp_usage_counters() + assert nlp_usage['asr_seconds_current_month'] == 4586 + assert nlp_usage['asr_seconds_all_time'] == 4728 + assert nlp_usage['mt_characters_current_month'] == 5473 + assert nlp_usage['mt_characters_all_time'] == 6726 + + def test_storage_usage(self): + calculator = UsageCalculator(self.anotheruser, None) + assert calculator.get_storage_usage() == 5 * self.expected_file_size() + + def test_submission_counters(self): + calculator = UsageCalculator(self.anotheruser, None) + submission_counters = calculator.get_submission_counters() + assert submission_counters['current_month'] == 5 + assert submission_counters['all_time'] == 5 + + def test_no_data(self): + calculator = UsageCalculator(self.someuser, None) + nlp_usage = calculator.get_nlp_usage_counters() + submission_counters = calculator.get_submission_counters() + + assert nlp_usage['asr_seconds_current_month'] == 0 + assert nlp_usage['asr_seconds_all_time'] == 0 + assert nlp_usage['mt_characters_current_month'] == 0 + assert nlp_usage['mt_characters_all_time'] == 0 + assert calculator.get_storage_usage() == 0 + assert submission_counters['current_month'] == 0 + assert submission_counters['all_time'] == 0 + + @override_settings(STRIPE_ENABLED=True) + def test_organization_setup(self): + organization = baker.make(Organization, id='org_abcd1234') + organization.add_user(user=self.anotheruser, is_admin=True) + organization.add_user(user=self.someuser, is_admin=True) + generate_enterprise_subscription(organization) + + calculator = UsageCalculator(self.someuser, organization) + submission_counters = calculator.get_submission_counters() + assert submission_counters['current_month'] == 5 + assert submission_counters['all_time'] == 5 + + nlp_usage = calculator.get_nlp_usage_counters() + assert nlp_usage['asr_seconds_current_month'] == 4586 + assert nlp_usage['asr_seconds_all_time'] == 4728 + assert nlp_usage['mt_characters_current_month'] == 5473 + assert nlp_usage['mt_characters_all_time'] == 6726 + + assert calculator.get_storage_usage() == 5 * self.expected_file_size() diff --git a/kpi/utils/usage_calculator.py b/kpi/utils/usage_calculator.py new file mode 100644 index 0000000000..1119babbe3 --- /dev/null +++ b/kpi/utils/usage_calculator.py @@ -0,0 +1,146 @@ +from typing import Optional + +from django.conf import settings +from django.db.models import Sum, Q +from django.db.models.functions import Coalesce +from django.utils import timezone + +from kobo.apps.kobo_auth.shortcuts import User +from kobo.apps.organizations.models import Organization +from kobo.apps.organizations.utils import ( + get_monthly_billing_dates, + get_yearly_billing_dates, +) +from kobo.apps.stripe.constants import ACTIVE_STRIPE_STATUSES +from kobo.apps.trackers.models import NLPUsageCounter +from kpi.deployment_backends.kc_access.shadow_models import ( + KobocatXForm, + KobocatDailyXFormSubmissionCounter, +) + + +class UsageCalculator: + def __init__(self, user: User, organization: Optional[Organization]): + self.user = user + self.organization = organization + + self._user_ids = [user.pk] + self._user_id_query = self._filter_by_user([user.pk]) + if organization and settings.STRIPE_ENABLED: + # if the user is in an organization and has an enterprise plan, get all org users + # we evaluate this queryset instead of using it as a subquery because it's referencing + # fields from the auth_user tables on kpi *and* kobocat, making getting results in a + # single query not feasible until those tables are combined + user_ids = list( + User.objects.filter( + organizations_organization__id=organization.id, + organizations_organization__djstripe_customers__subscriptions__status__in=ACTIVE_STRIPE_STATUSES, + organizations_organization__djstripe_customers__subscriptions__items__price__product__metadata__has_key='plan_type', + organizations_organization__djstripe_customers__subscriptions__items__price__product__metadata__plan_type='enterprise', + ).values_list('pk', flat=True)[ + : settings.ORGANIZATION_USER_LIMIT + ] + ) + if user_ids: + self._user_ids = user_ids + self._user_id_query = self._filter_by_user(user_ids) + + now = timezone.now() + self.current_month_start, self.current_month_end = ( + get_monthly_billing_dates(organization) + ) + self.current_year_start, self.current_year_end = ( + get_yearly_billing_dates(organization) + ) + self.current_month_filter = Q( + date__range=[self.current_month_start, now] + ) + self.current_year_filter = Q(date__range=[self.current_year_start, now]) + + def _filter_by_user(self, user_ids: list) -> Q: + """ + Turns a list of user ids into a query object to filter by + """ + return Q(user_id__in=user_ids) + + def get_nlp_usage_counters(self): + nlp_tracking = ( + NLPUsageCounter.objects.only( + 'date', 'total_asr_seconds', 'total_mt_characters' + ) + .filter(self._user_id_query) + .aggregate( + asr_seconds_current_year=Coalesce( + Sum('total_asr_seconds', filter=self.current_year_filter), 0 + ), + mt_characters_current_year=Coalesce( + Sum('total_mt_characters', filter=self.current_year_filter), + 0, + ), + asr_seconds_current_month=Coalesce( + Sum('total_asr_seconds', filter=self.current_month_filter), + 0, + ), + mt_characters_current_month=Coalesce( + Sum( + 'total_mt_characters', filter=self.current_month_filter + ), + 0, + ), + asr_seconds_all_time=Coalesce(Sum('total_asr_seconds'), 0), + mt_characters_all_time=Coalesce(Sum('total_mt_characters'), 0), + ) + ) + + total_nlp_usage = {} + for nlp_key, count in nlp_tracking.items(): + total_nlp_usage[nlp_key] = count if count is not None else 0 + + return total_nlp_usage + + def get_storage_usage(self): + """ + Get the storage used by non-(soft-)deleted projects for all users + + Users are represented by their ids with `self._user_ids` + """ + xforms = ( + KobocatXForm.objects.only('attachment_storage_bytes', 'id') + .exclude(pending_delete=True) + .filter(self._user_id_query) + ) + + total_storage_bytes = xforms.aggregate( + bytes_sum=Coalesce(Sum('attachment_storage_bytes'), 0), + ) + + return total_storage_bytes['bytes_sum'] or 0 + + def get_submission_counters(self): + """ + Calculate submissions for all users' projects even their deleted ones + + Users are represented by their ids with `self._user_ids` + """ + submission_count = ( + KobocatDailyXFormSubmissionCounter.objects.only( + 'counter', 'user_id' + ) + .filter(self._user_id_query) + .aggregate( + all_time=Coalesce(Sum('counter'), 0), + current_year=Coalesce( + Sum('counter', filter=self.current_year_filter), 0 + ), + current_month=Coalesce( + Sum('counter', filter=self.current_month_filter), 0 + ), + ) + ) + total_submission_count = {} + for submission_key, count in submission_count.items(): + total_submission_count[submission_key] = ( + count if count is not None else 0 + ) + + return total_submission_count From 6317be6bcb65172fdc10040ca646bce20884a5a1 Mon Sep 17 00:00:00 2001 From: Guillermo Date: Wed, 21 Aug 2024 21:36:36 -0600 Subject: [PATCH 04/10] Remove unused imports due to refactor --- kpi/serializers/v2/service_usage.py | 35 +++++--------- kpi/tests/api/v2/test_api_service_usage.py | 53 ++++++++-------------- kpi/utils/usage_calculator.py | 2 +- 3 files changed, 33 insertions(+), 57 deletions(-) diff --git a/kpi/serializers/v2/service_usage.py b/kpi/serializers/v2/service_usage.py index 822f8ffe51..93bc1236f8 100644 --- a/kpi/serializers/v2/service_usage.py +++ b/kpi/serializers/v2/service_usage.py @@ -1,22 +1,11 @@ -from django.conf import settings -from django.db.models import Sum, Q -from django.db.models.functions import Coalesce -from django.utils import timezone from rest_framework import serializers from rest_framework.fields import empty -from kobo.apps.kobo_auth.shortcuts import User from kobo.apps.organizations.models import Organization from kobo.apps.organizations.utils import ( get_monthly_billing_dates, get_yearly_billing_dates, ) -from kobo.apps.stripe.constants import ACTIVE_STRIPE_STATUSES -from kobo.apps.trackers.models import NLPUsageCounter -from kpi.deployment_backends.kc_access.shadow_models import ( - KobocatXForm, - KobocatDailyXFormSubmissionCounter, -) from kpi.deployment_backends.kobocat_backend import KobocatDeploymentBackend from kpi.models.asset import Asset from kpi.utils.usage_calculator import UsageCalculator @@ -130,23 +119,23 @@ def __init__(self, instance=None, data=empty, **kwargs): ).first() self.calculator = UsageCalculator(instance, organization) - def get_total_nlp_usage(self, user): - return self.calculator.get_nlp_usage_counters() - - def get_total_submission_count(self, user): - return self.calculator.get_submission_counters() - - def get_total_storage_bytes(self, user): - return self.calculator.get_storage_usage() + def get_current_month_end(self, user): + return self.calculator.current_month_end.isoformat() def get_current_month_start(self, user): return self.calculator.current_month_start.isoformat() - def get_current_month_end(self, user): - return self.calculator.current_month_end.isoformat() + def get_current_year_end(self, user): + return self.calculator.current_year_end.isoformat() def get_current_year_start(self, user): return self.calculator.current_year_start.isoformat() - def get_current_year_end(self, user): - return self.calculator.current_year_end.isoformat() + def get_total_nlp_usage(self, user): + return self.calculator.get_nlp_usage_counters() + + def get_total_submission_count(self, user): + return self.calculator.get_submission_counters() + + def get_total_storage_bytes(self, user): + return self.calculator.get_storage_usage() diff --git a/kpi/tests/api/v2/test_api_service_usage.py b/kpi/tests/api/v2/test_api_service_usage.py index af5deb90f8..c00f2adb92 100644 --- a/kpi/tests/api/v2/test_api_service_usage.py +++ b/kpi/tests/api/v2/test_api_service_usage.py @@ -1,19 +1,6 @@ -# coding: utf-8 -import os.path -import uuid - -from dateutil.relativedelta import relativedelta -from django.conf import settings from django.urls import reverse -from django.utils import timezone from rest_framework import status -from kobo.apps.kobo_auth.shortcuts import User -from kobo.apps.openrosa.apps.logger.models import ( - XForm, - DailyXFormSubmissionCounter, -) -from kobo.apps.trackers.models import NLPUsageCounter from kpi.models import Asset from kpi.tests.test_usage_calculator import BaseUsageCalculatorTestCase @@ -78,26 +65,6 @@ def test_multiple_forms(self): self.expected_file_size() * 3 ) - def test_service_usages_with_projects_in_trash_bin(self): - self.test_multiple_forms() - # Simulate trash bin - self.asset.pending_delete = True - self.asset.save( - update_fields=['pending_delete'], - create_version=False, - adjust_content=False, - ) - self.xform.pending_delete = True - self.xform.save(update_fields=['pending_delete']) - - # Retry endpoint - url = reverse(self._get_endpoint('service-usage-list')) - response = self.client.get(url) - - assert response.data['total_submission_count']['current_month'] == 3 - assert response.data['total_submission_count']['all_time'] == 3 - assert response.data['total_storage_bytes'] == 0 - def test_no_data(self): """ Test the endpoint functions when assets have no data @@ -143,3 +110,23 @@ def test_no_deployment(self): assert response.data['total_submission_count']['all_time'] == 0 assert response.data['total_nlp_usage']['asr_seconds_all_time'] == 0 assert response.data['total_storage_bytes'] == 0 + + def test_service_usages_with_projects_in_trash_bin(self): + self.test_multiple_forms() + # Simulate trash bin + self.asset.pending_delete = True + self.asset.save( + update_fields=['pending_delete'], + create_version=False, + adjust_content=False, + ) + self.xform.pending_delete = True + self.xform.save(update_fields=['pending_delete']) + + # Retry endpoint + url = reverse(self._get_endpoint('service-usage-list')) + response = self.client.get(url) + + assert response.data['total_submission_count']['current_month'] == 3 + assert response.data['total_submission_count']['all_time'] == 3 + assert response.data['total_storage_bytes'] == 0 diff --git a/kpi/utils/usage_calculator.py b/kpi/utils/usage_calculator.py index 1119babbe3..e9d8c845b1 100644 --- a/kpi/utils/usage_calculator.py +++ b/kpi/utils/usage_calculator.py @@ -14,8 +14,8 @@ from kobo.apps.stripe.constants import ACTIVE_STRIPE_STATUSES from kobo.apps.trackers.models import NLPUsageCounter from kpi.deployment_backends.kc_access.shadow_models import ( - KobocatXForm, KobocatDailyXFormSubmissionCounter, + KobocatXForm, ) From a711471e557c0207a55f875a81ecae9fb0ce8358 Mon Sep 17 00:00:00 2001 From: Guillermo Date: Thu, 22 Aug 2024 11:57:42 -0600 Subject: [PATCH 05/10] Change name to specify service usage in the class name Signed-off-by: Guillermo Signed-off-by: Guillermo --- .../apps/stripe/tests/test_organization_usage.py | 6 +++--- kpi/serializers/v2/service_usage.py | 4 ++-- kpi/tests/api/v2/test_api_service_usage.py | 4 ++-- kpi/tests/test_usage_calculator.py | 16 ++++++++-------- kpi/utils/usage_calculator.py | 2 +- 5 files changed, 16 insertions(+), 16 deletions(-) diff --git a/kobo/apps/stripe/tests/test_organization_usage.py b/kobo/apps/stripe/tests/test_organization_usage.py index 9a5af8a1ea..89d8ff5ba6 100644 --- a/kobo/apps/stripe/tests/test_organization_usage.py +++ b/kobo/apps/stripe/tests/test_organization_usage.py @@ -15,13 +15,13 @@ from kobo.apps.organizations.models import Organization, OrganizationUser from kobo.apps.stripe.tests.utils import generate_enterprise_subscription, generate_plan_subscription from kobo.apps.trackers.tests.submission_utils import create_mock_assets, add_mock_submissions -from kpi.tests.test_usage_calculator import BaseUsageCalculatorTestCase +from kpi.tests.test_usage_calculator import ServiceUsageBaseTestCase from kpi.tests.api.v2.test_api_asset_usage import AssetUsageAPITestCase from rest_framework import status -class OrganizationServiceUsageAPIMultiUserTestCase(BaseUsageCalculatorTestCase): +class OrganizationServiceUsageAPIMultiUserTestCase(ServiceUsageBaseTestCase): """ Test organization service usage when Stripe is enabled. @@ -149,7 +149,7 @@ def test_endpoint_is_cached(self): self.expected_file_size() * self.expected_submissions_multi ) -class OrganizationServiceUsageAPITestCase(BaseUsageCalculatorTestCase): +class OrganizationServiceUsageAPITestCase(ServiceUsageBaseTestCase): org_id = 'orgAKWMFskafsngf' @classmethod diff --git a/kpi/serializers/v2/service_usage.py b/kpi/serializers/v2/service_usage.py index 93bc1236f8..99d022ac44 100644 --- a/kpi/serializers/v2/service_usage.py +++ b/kpi/serializers/v2/service_usage.py @@ -8,7 +8,7 @@ ) from kpi.deployment_backends.kobocat_backend import KobocatDeploymentBackend from kpi.models.asset import Asset -from kpi.utils.usage_calculator import UsageCalculator +from kpi.utils.usage_calculator import ServiceUsageCalculator class AssetUsageSerializer(serializers.HyperlinkedModelSerializer): @@ -117,7 +117,7 @@ def __init__(self, instance=None, data=empty, **kwargs): organization_users__user_id=instance.id, id=organization_id, ).first() - self.calculator = UsageCalculator(instance, organization) + self.calculator = ServiceUsageCalculator(instance, organization) def get_current_month_end(self, user): return self.calculator.current_month_end.isoformat() diff --git a/kpi/tests/api/v2/test_api_service_usage.py b/kpi/tests/api/v2/test_api_service_usage.py index c00f2adb92..37ac4c54a2 100644 --- a/kpi/tests/api/v2/test_api_service_usage.py +++ b/kpi/tests/api/v2/test_api_service_usage.py @@ -2,10 +2,10 @@ from rest_framework import status from kpi.models import Asset -from kpi.tests.test_usage_calculator import BaseUsageCalculatorTestCase +from kpi.tests.test_usage_calculator import BaseServiceUsageTestCase -class ServiceUsageAPITestCase(BaseUsageCalculatorTestCase): +class ServiceUsageAPITestCase(BaseServiceUsageTestCase): def test_anonymous_user(self): """ Test that the endpoint is forbidden to anonymous user diff --git a/kpi/tests/test_usage_calculator.py b/kpi/tests/test_usage_calculator.py index 0cc6236e8e..49df95ec71 100644 --- a/kpi/tests/test_usage_calculator.py +++ b/kpi/tests/test_usage_calculator.py @@ -17,11 +17,11 @@ from kobo.apps.trackers.models import NLPUsageCounter from kpi.models import Asset from kpi.tests.base_test_case import BaseAssetTestCase -from kpi.utils.usage_calculator import UsageCalculator +from kpi.utils.usage_calculator import ServiceUsageCalculator from kpi.urls.router_api_v2 import URL_NAMESPACE as ROUTER_URL_NAMESPACE -class BaseUsageCalculatorTestCase(BaseAssetTestCase): +class BaseServiceUsageTestCase(BaseAssetTestCase): """ This class contains setup logic and utility functions to test usage calculations @@ -208,7 +208,7 @@ def update_xform_counters(self, asset: Asset, submissions: int = 0): self.counter.save() -class UsageCalculatorTestCase(BaseUsageCalculatorTestCase): +class ServiceUsageCalculatorTestCase(BaseServiceUsageTestCase): def setUp(self): super().setUp() self._create_asset() @@ -216,7 +216,7 @@ def setUp(self): self.add_submissions(count=5) def test_nlp_usage_counters(self): - calculator = UsageCalculator(self.anotheruser, None) + calculator = ServiceUsageCalculator(self.anotheruser, None) nlp_usage = calculator.get_nlp_usage_counters() assert nlp_usage['asr_seconds_current_month'] == 4586 assert nlp_usage['asr_seconds_all_time'] == 4728 @@ -224,17 +224,17 @@ def test_nlp_usage_counters(self): assert nlp_usage['mt_characters_all_time'] == 6726 def test_storage_usage(self): - calculator = UsageCalculator(self.anotheruser, None) + calculator = ServiceUsageCalculator(self.anotheruser, None) assert calculator.get_storage_usage() == 5 * self.expected_file_size() def test_submission_counters(self): - calculator = UsageCalculator(self.anotheruser, None) + calculator = ServiceUsageCalculator(self.anotheruser, None) submission_counters = calculator.get_submission_counters() assert submission_counters['current_month'] == 5 assert submission_counters['all_time'] == 5 def test_no_data(self): - calculator = UsageCalculator(self.someuser, None) + calculator = ServiceUsageCalculator(self.someuser, None) nlp_usage = calculator.get_nlp_usage_counters() submission_counters = calculator.get_submission_counters() @@ -253,7 +253,7 @@ def test_organization_setup(self): organization.add_user(user=self.someuser, is_admin=True) generate_enterprise_subscription(organization) - calculator = UsageCalculator(self.someuser, organization) + calculator = ServiceUsageCalculator(self.someuser, organization) submission_counters = calculator.get_submission_counters() assert submission_counters['current_month'] == 5 assert submission_counters['all_time'] == 5 diff --git a/kpi/utils/usage_calculator.py b/kpi/utils/usage_calculator.py index e9d8c845b1..481168ede9 100644 --- a/kpi/utils/usage_calculator.py +++ b/kpi/utils/usage_calculator.py @@ -19,7 +19,7 @@ ) -class UsageCalculator: +class ServiceUsageCalculator: def __init__(self, user: User, organization: Optional[Organization]): self.user = user self.organization = organization From 9e7a4c9a47d94d7c199082f52b5a8c3a24a80d1a Mon Sep 17 00:00:00 2001 From: Guillermo Date: Thu, 22 Aug 2024 11:57:57 -0600 Subject: [PATCH 06/10] Fix broken test related to service usage calculations --- .../tests/api/v2/test_api.py | 10 ++++----- kobo/apps/project_ownership/tests/utils.py | 21 +++++++++++-------- 2 files changed, 17 insertions(+), 14 deletions(-) diff --git a/kobo/apps/project_ownership/tests/api/v2/test_api.py b/kobo/apps/project_ownership/tests/api/v2/test_api.py index 9b600815f7..526270aa3b 100644 --- a/kobo/apps/project_ownership/tests/api/v2/test_api.py +++ b/kobo/apps/project_ownership/tests/api/v2/test_api.py @@ -16,7 +16,7 @@ InviteStatusChoices, Transfer, ) -from kobo.apps.project_ownership.tests.utils import MockServiceUsageSerializer +from kobo.apps.project_ownership.tests.utils import MockServiceUsageCalculator from kobo.apps.trackers.utils import update_nlp_counter from kpi.constants import PERM_VIEW_ASSET @@ -357,12 +357,12 @@ def __add_submissions(self): self.submissions = submissions @patch( - 'kpi.serializers.v2.service_usage.ServiceUsageSerializer._get_storage_usage', - new=MockServiceUsageSerializer._get_storage_usage + 'kpi.utils.usage_calculator.ServiceUsageCalculator.get_storage_usage', + new=MockServiceUsageCalculator.get_storage_usage ) @patch( - 'kpi.serializers.v2.service_usage.ServiceUsageSerializer._get_submission_counters', - new=MockServiceUsageSerializer._get_submission_counters + 'kpi.utils.usage_calculator.ServiceUsageCalculator.get_submission_counters', + new=MockServiceUsageCalculator.get_submission_counters ) @patch( 'kobo.apps.project_ownership.models.transfer.reset_kc_permissions', diff --git a/kobo/apps/project_ownership/tests/utils.py b/kobo/apps/project_ownership/tests/utils.py index fe46f81c9e..602f7e0775 100644 --- a/kobo/apps/project_ownership/tests/utils.py +++ b/kobo/apps/project_ownership/tests/utils.py @@ -2,24 +2,25 @@ from kpi.models.asset import Asset -class MockServiceUsageSerializer: +class MockServiceUsageCalculator: - def _get_storage_usage(self): + def get_storage_usage(self): assets = Asset.objects.annotate(user_id=F('owner_id')).filter( self._user_id_query ) - self._total_storage_bytes = 0 + total_storage_bytes = 0 for asset in assets: if asset.has_deployment: for submission in asset.deployment.get_submissions(asset.owner): - self._total_storage_bytes += sum( + total_storage_bytes += sum( [att['bytes'] for att in submission['_attachments']] ) + return total_storage_bytes - def _get_submission_counters(self, month_filter, year_filter): - self._total_submission_count = { + def get_submission_counters(self): + total_submission_count = { 'all_time': 0, 'current_year': 0, 'current_month': 0, @@ -30,6 +31,8 @@ def _get_submission_counters(self, month_filter, year_filter): for asset in assets: if asset.has_deployment: submissions = asset.deployment.get_submissions(asset.owner) - self._total_submission_count['all_time'] += len(submissions) - self._total_submission_count['current_year'] += len(submissions) - self._total_submission_count['current_month'] += len(submissions) + total_submission_count['all_time'] += len(submissions) + total_submission_count['current_year'] += len(submissions) + total_submission_count['current_month'] += len(submissions) + + return total_submission_count From 22ff2b70a82faba2d1fd741ffd4870d11a1f678b Mon Sep 17 00:00:00 2001 From: Guillermo Date: Thu, 22 Aug 2024 12:11:42 -0600 Subject: [PATCH 07/10] Fix broken test due to change in class name --- kobo/apps/stripe/tests/test_organization_usage.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/kobo/apps/stripe/tests/test_organization_usage.py b/kobo/apps/stripe/tests/test_organization_usage.py index 89d8ff5ba6..16da3d55f5 100644 --- a/kobo/apps/stripe/tests/test_organization_usage.py +++ b/kobo/apps/stripe/tests/test_organization_usage.py @@ -15,13 +15,13 @@ from kobo.apps.organizations.models import Organization, OrganizationUser from kobo.apps.stripe.tests.utils import generate_enterprise_subscription, generate_plan_subscription from kobo.apps.trackers.tests.submission_utils import create_mock_assets, add_mock_submissions -from kpi.tests.test_usage_calculator import ServiceUsageBaseTestCase +from kpi.tests.test_usage_calculator import BaseServiceUsageTestCase from kpi.tests.api.v2.test_api_asset_usage import AssetUsageAPITestCase from rest_framework import status -class OrganizationServiceUsageAPIMultiUserTestCase(ServiceUsageBaseTestCase): +class OrganizationServiceUsageAPIMultiUserTestCase(BaseServiceUsageTestCase): """ Test organization service usage when Stripe is enabled. @@ -149,7 +149,7 @@ def test_endpoint_is_cached(self): self.expected_file_size() * self.expected_submissions_multi ) -class OrganizationServiceUsageAPITestCase(ServiceUsageBaseTestCase): +class OrganizationServiceUsageAPITestCase(BaseServiceUsageTestCase): org_id = 'orgAKWMFskafsngf' @classmethod From 415b51fa3cf8c46ef600b698d8dd0b5e8a04ed69 Mon Sep 17 00:00:00 2001 From: Guillermo Date: Fri, 23 Aug 2024 12:34:51 -0600 Subject: [PATCH 08/10] Fix broken things after merge --- kpi/tests/api/v2/test_api_service_usage.py | 20 -------------------- kpi/tests/test_usage_calculator.py | 9 ++++++--- kpi/utils/usage_calculator.py | 6 +----- 3 files changed, 7 insertions(+), 28 deletions(-) diff --git a/kpi/tests/api/v2/test_api_service_usage.py b/kpi/tests/api/v2/test_api_service_usage.py index 5a0de08de0..68af6b2bbb 100644 --- a/kpi/tests/api/v2/test_api_service_usage.py +++ b/kpi/tests/api/v2/test_api_service_usage.py @@ -132,23 +132,3 @@ def test_no_deployment(self): assert response.data['total_submission_count']['all_time'] == 0 assert response.data['total_nlp_usage']['asr_seconds_all_time'] == 0 assert response.data['total_storage_bytes'] == 0 - - def test_service_usages_with_projects_in_trash_bin(self): - self.test_multiple_forms() - # Simulate trash bin - self.asset.pending_delete = True - self.asset.save( - update_fields=['pending_delete'], - create_version=False, - adjust_content=False, - ) - self.xform.pending_delete = True - self.xform.save(update_fields=['pending_delete']) - - # Retry endpoint - url = reverse(self._get_endpoint('service-usage-list')) - response = self.client.get(url) - - assert response.data['total_submission_count']['current_month'] == 3 - assert response.data['total_submission_count']['all_time'] == 3 - assert response.data['total_storage_bytes'] == 0 diff --git a/kpi/tests/test_usage_calculator.py b/kpi/tests/test_usage_calculator.py index 5b53f6f944..be4255a235 100644 --- a/kpi/tests/test_usage_calculator.py +++ b/kpi/tests/test_usage_calculator.py @@ -4,6 +4,7 @@ from dateutil.relativedelta import relativedelta from django.conf import settings from django.test import override_settings +from django.urls import reverse from django.utils import timezone from model_bakery import baker @@ -145,16 +146,18 @@ def add_submissions(self, count=2): self.attachment_id = self.attachment_id + 2 submissions.append(submission) - self.asset.deployment.mock_submissions(submissions, flush_db=False) + self.asset.deployment.mock_submissions(submissions) def expected_file_size(self): """ Calculate the expected combined file size for the test audio clip and image """ return os.path.getsize( - settings.BASE_DIR + '/kpi/tests/audio_conversion_test_clip.3gp' + settings.BASE_DIR + + '/kpi/fixtures/attachments/audio_conversion_test_clip.3gp' ) + os.path.getsize( - settings.BASE_DIR + '/kpi/tests/audio_conversion_test_image.jpg' + settings.BASE_DIR + + '/kpi/fixtures/attachments/audio_conversion_test_image.jpg' ) diff --git a/kpi/utils/usage_calculator.py b/kpi/utils/usage_calculator.py index 1344d61c94..71d3d24175 100644 --- a/kpi/utils/usage_calculator.py +++ b/kpi/utils/usage_calculator.py @@ -17,10 +17,6 @@ ) from kobo.apps.stripe.constants import ACTIVE_STRIPE_STATUSES from kobo.apps.trackers.models import NLPUsageCounter -from kpi.deployment_backends.kc_access.shadow_models import ( - KobocatDailyXFormSubmissionCounter, - KobocatXForm, -) class ServiceUsageCalculator: @@ -138,7 +134,7 @@ def get_submission_counters(self): total_submission_count = {} for submission_key, count in submission_count.items(): - self.total_submission_count[submission_key] = ( + total_submission_count[submission_key] = ( count if count is not None else 0 ) From 79ea9b0004a72530b3c7de9bd474c99a2ee64bd2 Mon Sep 17 00:00:00 2001 From: Guillermo Date: Fri, 23 Aug 2024 19:37:34 -0600 Subject: [PATCH 09/10] Fix test and remove obsolete class --- .../tests/api/v2/test_api.py | 9 ----- kobo/apps/project_ownership/tests/utils.py | 38 ------------------- 2 files changed, 47 deletions(-) delete mode 100644 kobo/apps/project_ownership/tests/utils.py diff --git a/kobo/apps/project_ownership/tests/api/v2/test_api.py b/kobo/apps/project_ownership/tests/api/v2/test_api.py index 9bd1e00a4d..59ec50bced 100644 --- a/kobo/apps/project_ownership/tests/api/v2/test_api.py +++ b/kobo/apps/project_ownership/tests/api/v2/test_api.py @@ -13,7 +13,6 @@ InviteStatusChoices, Transfer, ) -from kobo.apps.project_ownership.tests.utils import MockServiceUsageCalculator from kobo.apps.trackers.utils import update_nlp_counter from kpi.constants import PERM_VIEW_ASSET @@ -351,14 +350,6 @@ def __add_submissions(self): self.asset.deployment.mock_submissions(submissions) self.submissions = submissions - @patch( - 'kpi.utils.usage_calculator.ServiceUsageCalculator.get_storage_usage', - new=MockServiceUsageCalculator.get_storage_usage - ) - @patch( - 'kpi.utils.usage_calculator.ServiceUsageCalculator.get_submission_counters', - new=MockServiceUsageCalculator.get_submission_counters - ) @patch( 'kobo.apps.project_ownership.models.transfer.reset_kc_permissions', MagicMock() diff --git a/kobo/apps/project_ownership/tests/utils.py b/kobo/apps/project_ownership/tests/utils.py deleted file mode 100644 index 602f7e0775..0000000000 --- a/kobo/apps/project_ownership/tests/utils.py +++ /dev/null @@ -1,38 +0,0 @@ -from django.db.models import F -from kpi.models.asset import Asset - - -class MockServiceUsageCalculator: - - def get_storage_usage(self): - - assets = Asset.objects.annotate(user_id=F('owner_id')).filter( - self._user_id_query - ) - - total_storage_bytes = 0 - for asset in assets: - if asset.has_deployment: - for submission in asset.deployment.get_submissions(asset.owner): - total_storage_bytes += sum( - [att['bytes'] for att in submission['_attachments']] - ) - return total_storage_bytes - - def get_submission_counters(self): - total_submission_count = { - 'all_time': 0, - 'current_year': 0, - 'current_month': 0, - } - assets = Asset.objects.annotate(user_id=F('owner_id')).filter( - self._user_id_query - ) - for asset in assets: - if asset.has_deployment: - submissions = asset.deployment.get_submissions(asset.owner) - total_submission_count['all_time'] += len(submissions) - total_submission_count['current_year'] += len(submissions) - total_submission_count['current_month'] += len(submissions) - - return total_submission_count From 9e4bb41748d9970904939216fda751e62535943f Mon Sep 17 00:00:00 2001 From: Guillermo Date: Tue, 27 Aug 2024 08:52:01 -0600 Subject: [PATCH 10/10] Remove unused class properties --- kpi/tests/test_usage_calculator.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/kpi/tests/test_usage_calculator.py b/kpi/tests/test_usage_calculator.py index be4255a235..921ab99c96 100644 --- a/kpi/tests/test_usage_calculator.py +++ b/kpi/tests/test_usage_calculator.py @@ -31,9 +31,6 @@ class BaseServiceUsageTestCase(BaseAssetTestCase): URL_NAMESPACE = ROUTER_URL_NAMESPACE - attachment_id = 0 - counter = None - def setUp(self): super().setUp() self.client.login(username='anotheruser', password='anotheruser') @@ -128,13 +125,11 @@ def add_submissions(self, count=2): '_uuid': str(uuid.uuid4()), '_attachments': [ { - 'id': self.attachment_id, 'download_url': 'http://testserver/anotheruser/audio_conversion_test_clip.3gp', 'filename': 'anotheruser/audio_conversion_test_clip.3gp', 'mimetype': 'video/3gpp', }, { - 'id': self.attachment_id + 1, 'download_url': 'http://testserver/anotheruser/audio_conversion_test_image.jpg', 'filename': 'anotheruser/audio_conversion_test_image.jpg', 'mimetype': 'image/jpeg', @@ -143,7 +138,6 @@ def add_submissions(self, count=2): '_submitted_by': 'anotheruser', } # increment the attachment ID for each attachment created - self.attachment_id = self.attachment_id + 2 submissions.append(submission) self.asset.deployment.mock_submissions(submissions)