Fix logic for limits and usage, fix characters usage type error and i…

…mprove unit tests

kobo/apps/stripe/constants.py

@@ -22,7 +22,7 @@
 ORGANIZATION_USAGE_MAX_CACHE_AGE = timedelta(minutes=15)
 
 USAGE_LIMIT_MAP = {
-    'character': 'mt_character',
+    'character': 'mt_characters',
     'seconds': 'asr_seconds',
     'storage': 'storage_bytes',
     'submission': 'submission',

kobo/apps/stripe/tests/test_one_time_addons_api.py

@@ -1,3 +1,4 @@
+from ddt import ddt, data
 from django.urls import reverse
 from django.utils import timezone
 from djstripe.models import (
@@ -18,6 +19,7 @@
 from kpi.tests.kpi_test_case import BaseTestCase
 
 
+@ddt
 class OneTimeAddOnAPITestCase(BaseTestCase):
     fixtures = ['test_data']
 
@@ -179,7 +181,8 @@ def test_not_own_addon(self):
         assert response_get_list.status_code == status.HTTP_200_OK
         assert response_get_list.data['results'] == []
 
-    def _assert_get_user_totals(self, usage_type):
+    @data('character', 'seconds')
+    def test_get_user_totals(self, usage_type):
         limit = 2000
         quantity = 5
         usage_limit_key = f'{USAGE_LIMIT_MAP[usage_type]}_limit'
@@ -208,9 +211,3 @@ def _assert_get_user_totals(self, usage_type):
         )
         assert total_limit == limit * (quantity + 2)
         assert remaining == limit * 2
-
-    def test_get_user_totals_seconds(self):
-        self._assert_get_user_totals('seconds')
-
-    def test_get_user_totals_character(self):
-        self._assert_get_user_totals('character')

kobo/apps/stripe/utils.py

@@ -12,7 +12,7 @@ def get_default_add_on_limits():
     return {
         'submission_limit': 0,
         'asr_seconds_limit': 0,
-        'mt_character_limit': 0,
+        'mt_characters_limit': 0,
     }
 
 

kobo/apps/trackers/tests/test_utils.py

@@ -1,3 +1,4 @@
+from ddt import data, ddt
 from django.utils import timezone
 from djstripe.models import Charge, PaymentIntent, Price, Product
 from model_bakery import baker
@@ -14,6 +15,7 @@
 from kpi.tests.kpi_test_case import BaseTestCase
 
 
+@ddt
 class TrackersUtilitiesTestCase(BaseTestCase):
     fixtures = ['test_data']
 
@@ -33,9 +35,6 @@ def setUp(self):
         )
         self.asset.deploy(backend='mock', active=True)
 
-    def test_org_usage_seconds(self):
-        self._test_organization_usage_utils('seconds')
-
     def _create_product(self, product_metadata):
         product = baker.make(
             Product,
@@ -85,11 +84,12 @@ def _make_payment(
         charge.save()
         return charge
 
-    def _test_organization_usage_utils(self, usage_type):
+    @data('character', 'seconds')
+    def test_organization_usage_utils(self, usage_type):
         stripe_key = f'{USAGE_LIMIT_MAP_STRIPE[usage_type]}_limit'
         usage_key = f'{USAGE_LIMIT_MAP[usage_type]}_limit'
         sub_metadata = {
-            stripe_key: 1234,
+            stripe_key: 1000,
             'product_type': 'plan',
             'plan_type': 'enterprise',
         }
@@ -98,19 +98,25 @@ def _test_organization_usage_utils(self, usage_type):
         )
         addon_metadata = {
             'product_type': 'addon_onetime',
-            usage_key: 12345,
+            usage_key: 2000,
             'valid_tags': 'all',
         }
         product, price = self._create_product(addon_metadata)
         self._make_payment(price, subscription.customer, quantity=2)
 
-        expected_limit = 12345 * 2 + 1234
+        total_limit = 2000 * 2 + 1000
         remaining = get_organization_remaining_usage(self.organization, usage_type)
-        assert remaining == expected_limit
+        assert remaining == total_limit
 
         update_nlp_counter(
             USAGE_LIMIT_MAP[usage_type], 1000, self.someuser.id, self.asset.id
         )
 
         remaining = get_organization_remaining_usage(self.organization, usage_type)
-        assert remaining == expected_limit - 1000
+        assert remaining == total_limit - 1000
+
+        update_nlp_counter(
+            USAGE_LIMIT_MAP[usage_type], 1500, self.someuser.id, self.asset.id
+        )
+        remaining = get_organization_remaining_usage(self.organization, usage_type)
+        assert remaining == total_limit - 2500

kobo/apps/trackers/utils.py

@@ -95,7 +95,7 @@ def get_organization_remaining_usage(
         organization,
         usage_type,
     )
-    plan_remaining = max(0, plan_limit - usage) # if negative, they have 0 remaining
+    plan_remaining = max(0, plan_limit - usage)  # if negative, they have 0 remaining
     total_remaining = addon_remaining + plan_remaining
 
     return total_remaining
@@ -107,6 +107,8 @@ def handle_usage_increment(
     """
     Increment the given usage type for this organization by the given amount
     """
+    PlanAddOn = apps.get_model('stripe', 'PlanAddOn')
+
     plan_limit = get_organization_plan_limit(organization, usage_type)
     current_usage = get_organization_usage(organization, usage_type)
     if current_usage is None:
@@ -117,5 +119,5 @@ def handle_usage_increment(
             amount if current_usage >= plan_limit else new_total_usage - plan_limit
         )
         PlanAddOn.increment_add_ons_for_organization(
-            organization.id, usage_type, increment
+            organization, usage_type, increment
         )

kpi/utils/usage_calculator.py

@@ -74,7 +74,7 @@ def get_nlp_usage_by_type(self, usage_key: str) -> int:
 
         cached_usage = {
             'asr_seconds': nlp_usage[f'asr_seconds_current_{interval}'],
-            'mt_character': nlp_usage[f'mt_characters_current_{interval}'],
+            'mt_characters': nlp_usage[f'mt_characters_current_{interval}'],
         }
 
         return cached_usage[usage_key]