From 59f492eccfc0138cb1f25fc075b060d92342c0a1 Mon Sep 17 00:00:00 2001 From: Guillermo Date: Wed, 2 Oct 2024 17:34:27 -0600 Subject: [PATCH] Fix logic for limits and usage, fix characters usage type error and improve unit tests --- kobo/apps/stripe/constants.py | 2 +- .../stripe/tests/test_one_time_addons_api.py | 11 ++++----- kobo/apps/stripe/utils.py | 2 +- kobo/apps/trackers/tests/test_utils.py | 24 ++++++++++++------- kobo/apps/trackers/utils.py | 6 +++-- kpi/utils/usage_calculator.py | 2 +- 6 files changed, 26 insertions(+), 21 deletions(-) diff --git a/kobo/apps/stripe/constants.py b/kobo/apps/stripe/constants.py index 46c3d7a8d1..fc8e573a0c 100644 --- a/kobo/apps/stripe/constants.py +++ b/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', diff --git a/kobo/apps/stripe/tests/test_one_time_addons_api.py b/kobo/apps/stripe/tests/test_one_time_addons_api.py index de7b02d53a..8ac192ddbb 100644 --- a/kobo/apps/stripe/tests/test_one_time_addons_api.py +++ b/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') diff --git a/kobo/apps/stripe/utils.py b/kobo/apps/stripe/utils.py index 7ff13a7f15..6fc3d657b7 100644 --- a/kobo/apps/stripe/utils.py +++ b/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, } diff --git a/kobo/apps/trackers/tests/test_utils.py b/kobo/apps/trackers/tests/test_utils.py index fecf8058e6..f0a640a7c3 100644 --- a/kobo/apps/trackers/tests/test_utils.py +++ b/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 diff --git a/kobo/apps/trackers/utils.py b/kobo/apps/trackers/utils.py index 5f99299294..f59a17812c 100644 --- a/kobo/apps/trackers/utils.py +++ b/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 ) diff --git a/kpi/utils/usage_calculator.py b/kpi/utils/usage_calculator.py index cce2f83294..f2b90a8a07 100644 --- a/kpi/utils/usage_calculator.py +++ b/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]