Merge branch 'beta-refactored' of github.com:kobotoolbox/kpi into 862…

…-reject_duplicate_submissions

kobo/apps/project_ownership/tests/api/v2/test_api.py

@@ -13,7 +13,6 @@
     InviteStatusChoices,
     Transfer,
 )
-from kobo.apps.project_ownership.tests.utils import MockServiceUsageSerializer
 from kobo.apps.trackers.utils import update_nlp_counter
 
 from kpi.constants import PERM_VIEW_ASSET

kobo/apps/stripe/tests/test_organization_usage.py

@@ -14,20 +14,20 @@
 
 from kobo.apps.kobo_auth.shortcuts import User
 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 kobo.apps.stripe.tests.utils import (
+    generate_enterprise_subscription,
+    generate_plan_subscription,
+)
+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(ServiceUsageAPIBase):
+class OrganizationServiceUsageAPIMultiUserTestCase(BaseServiceUsageTestCase):
     """
     Test organization service usage when Stripe is enabled.
 
@@ -162,7 +162,7 @@ def test_endpoint_is_cached(self):
         )
 
 
-class OrganizationServiceUsageAPITestCase(ServiceUsageAPIBase):
+class OrganizationServiceUsageAPITestCase(BaseServiceUsageTestCase):
     org_id = 'orgAKWMFskafsngf'
 
     @classmethod

kpi/serializers/v2/service_usage.py

@@ -1,21 +1,14 @@
-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.openrosa.apps.logger.models import (
-    DailyXFormSubmissionCounter,
-    XForm,
-)
 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 kobo.apps.organizations.utils import (
+    get_monthly_billing_dates,
+    get_yearly_billing_dates,
+)
 from kpi.deployment_backends.openrosa_backend import OpenRosaDeploymentBackend
 from kpi.models.asset import Asset
+from kpi.utils.usage_calculator import ServiceUsageCalculator
 
 
 class AssetUsageSerializer(serializers.HyperlinkedModelSerializer):
@@ -117,158 +110,32 @@ class ServiceUsageSerializer(serializers.Serializer):
 
     def __init__(self, instance=None, data=empty, **kwargs):
         super().__init__(instance=instance, data=data, **kwargs)
+        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 = ServiceUsageCalculator(instance, organization)
 
-        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)
-
-    def get_total_nlp_usage(self, user):
-        return self._total_nlp_usage
-
-    def get_total_submission_count(self, user):
-        return self._total_submission_count
-
-    def get_total_storage_bytes(self, user):
-        return self._total_storage_bytes
-
-    def get_current_month_start(self, user):
-        return self._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()
+    def get_current_month_start(self, user):
+        return self.calculator.current_month_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
+        return self.calculator.current_year_end.isoformat()
 
-        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 = XForm.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_current_year_start(self, user):
+        return self.calculator.current_year_start.isoformat()
 
-    def _get_submission_counters(self, month_filter, year_filter):
-        """
-        Calculate submissions for all users' projects even their deleted ones
+    def get_total_nlp_usage(self, user):
+        return self.calculator.get_nlp_usage_counters()
 
-        Users are represented by their ids with `self._user_ids`
-        """
-        submission_count = DailyXFormSubmissionCounter.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
-            ),
-        )
+    def get_total_submission_count(self, user):
+        return self.calculator.get_submission_counters()
 
-        for submission_key, count in submission_count.items():
-            self._total_submission_count[submission_key] = (
-                count if count is not None else 0
-            )
+    def get_total_storage_bytes(self, user):
+        return self.calculator.get_storage_usage()