From 7d8e8a2286777f5ed235b306ff2e41538f7d551c Mon Sep 17 00:00:00 2001 From: Zach Aysan Date: Mon, 8 Apr 2024 17:56:18 +0000 Subject: [PATCH 01/22] Fix naming and add api usage model --- api/organisations/models.py | 14 +++++++++++++- api/organisations/views.py | 6 +++--- .../test_unit_organisations_models.py | 10 +++++----- .../organisations/test_unit_organisations_views.py | 12 ++++++------ 4 files changed, 27 insertions(+), 15 deletions(-) diff --git a/api/organisations/models.py b/api/organisations/models.py index bb377ae1ea07..f306a825fddb 100644 --- a/api/organisations/models.py +++ b/api/organisations/models.py @@ -455,7 +455,7 @@ def erase_api_notifications(self): self.organisation.api_usage_notifications.all().delete() -class OranisationAPIUsageNotification(models.Model): +class OrganisationAPIUsageNotification(models.Model): organisation = models.ForeignKey( Organisation, on_delete=models.CASCADE, related_name="api_usage_notifications" ) @@ -478,3 +478,15 @@ class HubspotOrganisation(models.Model): hubspot_id = models.CharField(max_length=100) created_at = models.DateTimeField(auto_now_add=True) updated_at = models.DateTimeField(auto_now=True) + + +class OrganisationAPIBilling(models.Model): + organisation = models.ForeignKey( + Organisation, on_delete=models.CASCADE, related_name="api_billing" + ) + api_overage = models.IntegerField(null=False) + immediate_invoice = models.BooleanField(null=False, default=False) + billed_at = models.DateTimeField(null=False) + + created_at = models.DateTimeField(null=True, auto_now_add=True) + updated_at = models.DateTimeField(null=True, auto_now=True) diff --git a/api/organisations/views.py b/api/organisations/views.py index bd3faf9e2efe..4c2e4c31f441 100644 --- a/api/organisations/views.py +++ b/api/organisations/views.py @@ -24,8 +24,8 @@ from organisations.chargebee import webhook_event_types, webhook_handlers from organisations.exceptions import OrganisationHasNoPaidSubscription from organisations.models import ( - OranisationAPIUsageNotification, Organisation, + OrganisationAPIUsageNotification, OrganisationRole, OrganisationWebhook, ) @@ -319,7 +319,7 @@ class OrganisationAPIUsageNotificationView(ListAPIView): def get_queryset(self): organisation = Organisation.objects.get(id=self.kwargs["organisation_pk"]) if not hasattr(organisation, "subscription_information_cache"): - return OranisationAPIUsageNotification.objects.none() + return OrganisationAPIUsageNotification.objects.none() subscription_cache = organisation.subscription_information_cache billing_starts_at = subscription_cache.current_billing_term_starts_at now = timezone.now() @@ -327,7 +327,7 @@ def get_queryset(self): month_delta = relativedelta(now, billing_starts_at).months period_starts_at = relativedelta(months=month_delta) + billing_starts_at - queryset = OranisationAPIUsageNotification.objects.filter( + queryset = OrganisationAPIUsageNotification.objects.filter( organisation_id=organisation.id, notified_at__gt=period_starts_at, ) diff --git a/api/tests/unit/organisations/test_unit_organisations_models.py b/api/tests/unit/organisations/test_unit_organisations_models.py index a19b4ad85fac..15344b32c70e 100644 --- a/api/tests/unit/organisations/test_unit_organisations_models.py +++ b/api/tests/unit/organisations/test_unit_organisations_models.py @@ -10,8 +10,8 @@ from environments.models import Environment from organisations.chargebee.metadata import ChargebeeObjMetadata from organisations.models import ( - OranisationAPIUsageNotification, Organisation, + OrganisationAPIUsageNotification, OrganisationSubscriptionInformationCache, Subscription, ) @@ -551,7 +551,7 @@ def test_reset_of_api_notifications(organisation: Organisation) -> None: ) # Create a notification which should be deleted shortly. - OranisationAPIUsageNotification.objects.create( + OrganisationAPIUsageNotification.objects.create( organisation=organisation, percent_usage=90, notified_at=now, @@ -559,7 +559,7 @@ def test_reset_of_api_notifications(organisation: Organisation) -> None: # Keep a notification which should not be deleted. organisation2 = Organisation.objects.create(name="Test org2") - oapiun = OranisationAPIUsageNotification.objects.create( + oapiun = OrganisationAPIUsageNotification.objects.create( organisation=organisation2, percent_usage=90, notified_at=now, @@ -570,5 +570,5 @@ def test_reset_of_api_notifications(organisation: Organisation) -> None: osic.save() # Then - assert OranisationAPIUsageNotification.objects.count() == 1 - assert OranisationAPIUsageNotification.objects.first() == oapiun + assert OrganisationAPIUsageNotification.objects.count() == 1 + assert OrganisationAPIUsageNotification.objects.first() == oapiun diff --git a/api/tests/unit/organisations/test_unit_organisations_views.py b/api/tests/unit/organisations/test_unit_organisations_views.py index 87c9a27729cb..275a9779c847 100644 --- a/api/tests/unit/organisations/test_unit_organisations_views.py +++ b/api/tests/unit/organisations/test_unit_organisations_views.py @@ -25,8 +25,8 @@ from organisations.chargebee.metadata import ChargebeeObjMetadata from organisations.invites.models import Invite from organisations.models import ( - OranisationAPIUsageNotification, Organisation, + OrganisationAPIUsageNotification, OrganisationRole, OrganisationSubscriptionInformationCache, OrganisationWebhook, @@ -1693,7 +1693,7 @@ def test_defaults_to_empty_api_notifications_when_no_subscription_information_ca ) now = timezone.now() - OranisationAPIUsageNotification.objects.create( + OrganisationAPIUsageNotification.objects.create( organisation=organisation, percent_usage=90, notified_at=now, @@ -1734,17 +1734,17 @@ def test_retrieves_api_usage_notifications( ) # Add three notifications, but we only get the 100% one. - OranisationAPIUsageNotification.objects.create( + OrganisationAPIUsageNotification.objects.create( organisation=organisation, percent_usage=90, notified_at=now, ) - OranisationAPIUsageNotification.objects.create( + OrganisationAPIUsageNotification.objects.create( organisation=organisation, percent_usage=75, notified_at=now, ) - OranisationAPIUsageNotification.objects.create( + OrganisationAPIUsageNotification.objects.create( organisation=organisation, percent_usage=100, notified_at=now, @@ -1785,7 +1785,7 @@ def test_doesnt_retrieve_stale_api_usage_notifications( ) # Create a notification in the past which should not be shown. - OranisationAPIUsageNotification.objects.create( + OrganisationAPIUsageNotification.objects.create( organisation=organisation, percent_usage=90, notified_at=now - timedelta(20), From 2e3b185c67e973afd2711ba09e703dfa3ecaf521 Mon Sep 17 00:00:00 2001 From: Zach Aysan Date: Mon, 8 Apr 2024 17:57:33 +0000 Subject: [PATCH 02/22] Add task to bill monthly users --- api/organisations/tasks.py | 89 ++++++++++++++++++++++++++++++++++++-- 1 file changed, 85 insertions(+), 4 deletions(-) diff --git a/api/organisations/tasks.py b/api/organisations/tasks.py index 279894ec3b2d..0021daa2cd6a 100644 --- a/api/organisations/tasks.py +++ b/api/organisations/tasks.py @@ -5,14 +5,20 @@ from dateutil.relativedelta import relativedelta from django.conf import settings from django.core.mail import send_mail +from django.db.models import F from django.template.loader import render_to_string from django.utils import timezone from integrations.flagsmith.client import get_client from organisations import subscription_info_cache +from organisations.chargebee import ( + add_1000_api_calls_scale_up, + add_1000_api_calls_start_up, +) from organisations.models import ( - OranisationAPIUsageNotification, Organisation, + OrganisationAPIBilling, + OrganisationAPIUsageNotification, OrganisationRole, Subscription, ) @@ -30,7 +36,13 @@ ALERT_EMAIL_SUBJECT, API_USAGE_ALERT_THRESHOLDS, ) -from .subscriptions.constants import SubscriptionCacheEntity +from .subscriptions.constants import ( + SCALE_UP, + SCALE_UP_V2, + STARTUP, + STARTUP_V2, + SubscriptionCacheEntity, +) logger = logging.getLogger(__name__) @@ -126,7 +138,7 @@ def send_admin_api_usage_notification( fail_silently=True, ) - OranisationAPIUsageNotification.objects.create( + OrganisationAPIUsageNotification.objects.create( organisation=organisation, percent_usage=matched_threshold, notified_at=timezone.now(), @@ -154,7 +166,7 @@ def _handle_api_usage_notifications(organisation: Organisation): matched_threshold = threshold - if OranisationAPIUsageNotification.objects.filter( + if OrganisationAPIUsageNotification.objects.filter( notified_at__gt=period_starts_at, percent_usage=matched_threshold, ).exists(): @@ -189,7 +201,76 @@ def handle_api_usage_notifications(): ) +def charge_for_api_call_count_overages(): + now = timezone.now() + api_usage_notified_at = now - timedelta(days=30) + month_window_start = timedelta(days=25) + month_window_end = timedelta(days=35) + closing_billing_term = now + timedelta(hours=12) + organisation_ids = set( + OrganisationAPIUsageNotification.objects.filter( + notified_at__gte=api_usage_notified_at, + percent_usage__gte=100, + ) + .exclude( + organisation__api_billing__billed_at__gt=api_usage_notified_at, + ) + .values_list("organisation_id", flat=True) + ) + + for organisation in Organisation.objects.filter( + id__in=organisation_ids, + subscription_information_cache__current_billing_term_ends_at__lte=closing_billing_term, + subscription_information_cache__current_billing_term_ends_at__gte=now, + subscription_information_cache__current_billing_term_starts_at__lte=F( + "subscription_information_cache__current_billing_term_ends_at" + ) + - month_window_start, + subscription_information_cache__current_billing_term_starts_at__gte=F( + "subscription_information_cache__current_billing_term_ends_at" + ) + - month_window_end, + ).select_related( + "subscription_information_cache", + ): + subscription_cache = organisation.subscription_information_cache + api_usage = get_current_api_usage(organisation.id, "30d") + api_usage_ratio = api_usage / subscription_cache.allowed_30d_api_calls + + if api_usage_ratio < 1.0: + logger.warning("API Usage does not match API Notification") + continue + + api_overage = api_usage - subscription_cache.allowed_30d_api_calls + + if organisation.subscription.plan in {SCALE_UP, SCALE_UP_V2}: + add_1000_api_calls_scale_up( + organisation.subscription.subscription_id, api_overage // 1000 + ) + elif organisation.subscription.plan in {STARTUP, STARTUP_V2}: + add_1000_api_calls_start_up( + organisation.subscription.subscription_id, api_overage // 1000 + ) + else: + logger.error( + f"Unable to bill for API overages for plan `{organisation.subscription.plan}`" + ) + continue + + # Save a copy of what was just billed in order to avoid + # double billing on a subsequent task run. + OrganisationAPIBilling.objects.create( + organisation=organisation, + api_overage=(1000 * (api_overage // 1000)), + immediate_invoice=False, + billed_at=now, + ) + + if settings.ENABLE_API_USAGE_ALERTING: register_recurring_task( run_every=timedelta(hours=12), )(handle_api_usage_notifications) + register_recurring_task( + run_every=timedelta(hours=12), + )(charge_for_api_call_count_overages) From 5034e38d30c44244bb51210958604aee0fd70524 Mon Sep 17 00:00:00 2001 From: Zach Aysan Date: Mon, 8 Apr 2024 17:58:08 +0000 Subject: [PATCH 03/22] Add api call billing to __init__.py --- api/organisations/chargebee/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/api/organisations/chargebee/__init__.py b/api/organisations/chargebee/__init__.py index 663a9b137aa2..cf03c4f2a052 100644 --- a/api/organisations/chargebee/__init__.py +++ b/api/organisations/chargebee/__init__.py @@ -1,4 +1,6 @@ from .chargebee import ( # noqa + add_1000_api_calls_scale_up, + add_1000_api_calls_start_up, add_single_seat, extract_subscription_metadata, get_customer_id_from_subscription_id, From 976e2ce6d3949ec4e79ec38d6fd93905e4d3f717 Mon Sep 17 00:00:00 2001 From: Zach Aysan Date: Mon, 8 Apr 2024 17:58:32 +0000 Subject: [PATCH 04/22] Add migration to fix naming and add api billing --- .../migrations/0053_create_api_billing.py | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) create mode 100644 api/organisations/migrations/0053_create_api_billing.py diff --git a/api/organisations/migrations/0053_create_api_billing.py b/api/organisations/migrations/0053_create_api_billing.py new file mode 100644 index 000000000000..49e2e1b19048 --- /dev/null +++ b/api/organisations/migrations/0053_create_api_billing.py @@ -0,0 +1,30 @@ +# Generated by Django 3.2.25 on 2024-04-08 15:07 + +from django.db import migrations, models +import django.db.models.deletion + + +class Migration(migrations.Migration): + + dependencies = [ + ('organisations', '0052_create_hubspot_organisation'), + ] + + operations = [ + migrations.RenameModel( + old_name='OranisationAPIUsageNotification', + new_name='OrganisationAPIUsageNotification', + ), + migrations.CreateModel( + name='OrganisationAPIBilling', + fields=[ + ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), + ('api_overage', models.IntegerField()), + ('immediate_invoice', models.BooleanField(default=False)), + ('billed_at', models.DateTimeField()), + ('created_at', models.DateTimeField(auto_now_add=True, null=True)), + ('updated_at', models.DateTimeField(auto_now=True, null=True)), + ('organisation', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='api_billing', to='organisations.organisation')), + ], + ), + ] From fb2652e9a5e94500e95936f1f1983702ad2d2ccf Mon Sep 17 00:00:00 2001 From: Zach Aysan Date: Mon, 8 Apr 2024 18:01:32 +0000 Subject: [PATCH 05/22] Create tasks for billing API usage --- .../test_unit_organisations_tasks.py | 249 +++++++++++++++++- 1 file changed, 237 insertions(+), 12 deletions(-) diff --git a/api/tests/unit/organisations/test_unit_organisations_tasks.py b/api/tests/unit/organisations/test_unit_organisations_tasks.py index 437e9841ce55..11ca3c546e97 100644 --- a/api/tests/unit/organisations/test_unit_organisations_tasks.py +++ b/api/tests/unit/organisations/test_unit_organisations_tasks.py @@ -9,8 +9,9 @@ from organisations.chargebee.metadata import ChargebeeObjMetadata from organisations.models import ( - OranisationAPIUsageNotification, Organisation, + OrganisationAPIBilling, + OrganisationAPIUsageNotification, OrganisationRole, OrganisationSubscriptionInformationCache, UserOrganisation, @@ -24,6 +25,7 @@ from organisations.tasks import ( ALERT_EMAIL_MESSAGE, ALERT_EMAIL_SUBJECT, + charge_for_api_call_count_overages, finish_subscription_cancellation, handle_api_usage_notifications, send_org_over_limit_alert, @@ -261,7 +263,7 @@ def test_handle_api_usage_notifications_when_feature_flag_is_off( assert len(mailoutbox) == 0 assert ( - OranisationAPIUsageNotification.objects.filter( + OrganisationAPIUsageNotification.objects.filter( organisation=organisation, ).count() == 0 @@ -294,7 +296,7 @@ def test_handle_api_usage_notifications_below_100( get_client_mock.return_value = client_mock client_mock.get_identity_flags.return_value.is_feature_enabled.return_value = True - assert not OranisationAPIUsageNotification.objects.filter( + assert not OrganisationAPIUsageNotification.objects.filter( organisation=organisation, ).exists() @@ -337,12 +339,12 @@ def test_handle_api_usage_notifications_below_100( assert email.to == ["admin@example.com"] assert ( - OranisationAPIUsageNotification.objects.filter( + OrganisationAPIUsageNotification.objects.filter( organisation=organisation, ).count() == 1 ) - api_usage_notification = OranisationAPIUsageNotification.objects.filter( + api_usage_notification = OrganisationAPIUsageNotification.objects.filter( organisation=organisation, ).first() @@ -352,12 +354,12 @@ def test_handle_api_usage_notifications_below_100( handle_api_usage_notifications() assert ( - OranisationAPIUsageNotification.objects.filter( + OrganisationAPIUsageNotification.objects.filter( organisation=organisation, ).count() == 1 ) - assert OranisationAPIUsageNotification.objects.first() == api_usage_notification + assert OrganisationAPIUsageNotification.objects.first() == api_usage_notification @pytest.mark.freeze_time("2023-01-19T09:09:47.325132+00:00") @@ -387,7 +389,7 @@ def test_handle_api_usage_notifications_above_100( get_client_mock.return_value = client_mock client_mock.get_identity_flags.return_value.is_feature_enabled.return_value = True - assert not OranisationAPIUsageNotification.objects.filter( + assert not OrganisationAPIUsageNotification.objects.filter( organisation=organisation, ).exists() @@ -431,12 +433,12 @@ def test_handle_api_usage_notifications_above_100( assert email.to == ["admin@example.com", "staff@example.com"] assert ( - OranisationAPIUsageNotification.objects.filter( + OrganisationAPIUsageNotification.objects.filter( organisation=organisation, ).count() == 1 ) - api_usage_notification = OranisationAPIUsageNotification.objects.filter( + api_usage_notification = OrganisationAPIUsageNotification.objects.filter( organisation=organisation, ).first() @@ -446,9 +448,232 @@ def test_handle_api_usage_notifications_above_100( handle_api_usage_notifications() assert ( - OranisationAPIUsageNotification.objects.filter( + OrganisationAPIUsageNotification.objects.filter( organisation=organisation, ).count() == 1 ) - assert OranisationAPIUsageNotification.objects.first() == api_usage_notification + assert OrganisationAPIUsageNotification.objects.first() == api_usage_notification + + +@pytest.mark.freeze_time("2023-01-19T09:09:47.325132+00:00") +def test_charge_for_api_call_count_overages_scale_up( + organisation: Organisation, + mocker: MockerFixture, +) -> None: + # Given + now = timezone.now() + OrganisationSubscriptionInformationCache.objects.create( + organisation=organisation, + allowed_seats=10, + allowed_projects=3, + allowed_30d_api_calls=10_000, + chargebee_email="test@example.com", + current_billing_term_starts_at=now - timedelta(days=30), + current_billing_term_ends_at=now + timedelta(hours=6), + ) + organisation.subscription.subscription_id = "fancy_sub_id23" + organisation.subscription.plan = "scale-up-v2" + organisation.subscription.save() + OrganisationAPIUsageNotification.objects.create( + organisation=organisation, + percent_usage=100, + notified_at=now, + ) + + mocker.patch("organisations.chargebee.chargebee.chargebee.Subscription.retrieve") + mock_chargebee_update = mocker.patch( + "organisations.chargebee.chargebee.chargebee.Subscription.update" + ) + + mock_api_usage = mocker.patch( + "organisations.tasks.get_current_api_usage", + ) + mock_api_usage.return_value = 12_005 + assert OrganisationAPIBilling.objects.count() == 0 + + # When + charge_for_api_call_count_overages() + + # Then + mock_chargebee_update.assert_called_once_with( + organisation.subscription.subscription_id, + { + "addons": [ + { + "id": "additional-api-scale-up-monthly", + "quantity": 2, # Two thousand API requests. + } + ], + "prorate": True, + "invoice_immediately": False, + }, + ) + + assert OrganisationAPIBilling.objects.count() == 1 + api_billing = OrganisationAPIBilling.objects.first() + assert api_billing.organisation == organisation + assert api_billing.api_overage == 2000 + assert api_billing.immediate_invoice is False + assert api_billing.billed_at == now + + +@pytest.mark.freeze_time("2023-01-19T09:09:47.325132+00:00") +def test_charge_for_api_call_count_overages_start_up( + organisation: Organisation, + mocker: MockerFixture, +) -> None: + # Given + now = timezone.now() + OrganisationSubscriptionInformationCache.objects.create( + organisation=organisation, + allowed_seats=10, + allowed_projects=3, + allowed_30d_api_calls=10_000, + chargebee_email="test@example.com", + current_billing_term_starts_at=now - timedelta(days=30), + current_billing_term_ends_at=now + timedelta(hours=6), + ) + organisation.subscription.subscription_id = "fancy_sub_id23" + organisation.subscription.plan = "startup-v2" + organisation.subscription.save() + OrganisationAPIUsageNotification.objects.create( + organisation=organisation, + percent_usage=100, + notified_at=now, + ) + + mocker.patch("organisations.chargebee.chargebee.chargebee.Subscription.retrieve") + mock_chargebee_update = mocker.patch( + "organisations.chargebee.chargebee.chargebee.Subscription.update" + ) + + mock_api_usage = mocker.patch( + "organisations.tasks.get_current_api_usage", + ) + mock_api_usage.return_value = 12_005 + assert OrganisationAPIBilling.objects.count() == 0 + + # When + charge_for_api_call_count_overages() + + # Then + mock_chargebee_update.assert_called_once_with( + organisation.subscription.subscription_id, + { + "addons": [ + { + "id": "additional-api-start-up-monthly", + "quantity": 2, # Two thousand API requests. + } + ], + "prorate": True, + "invoice_immediately": False, + }, + ) + + assert OrganisationAPIBilling.objects.count() == 1 + api_billing = OrganisationAPIBilling.objects.first() + assert api_billing.organisation == organisation + assert api_billing.api_overage == 2000 + assert api_billing.immediate_invoice is False + assert api_billing.billed_at == now + + # Now attempt to rebill the account should fail + calls_mock = mocker.patch( + "organisations.tasks.add_1000_api_calls_start_up", + ) + charge_for_api_call_count_overages() + assert OrganisationAPIBilling.objects.count() == 1 + calls_mock.assert_not_called() + + +@pytest.mark.freeze_time("2023-01-19T09:09:47.325132+00:00") +def test_charge_for_api_call_count_overages_with_yearly_account( + organisation: Organisation, + mocker: MockerFixture, +) -> None: + # Given + now = timezone.now() + OrganisationSubscriptionInformationCache.objects.create( + organisation=organisation, + allowed_seats=10, + allowed_projects=3, + allowed_30d_api_calls=10_000, + chargebee_email="test@example.com", + current_billing_term_starts_at=now - timedelta(days=365), + current_billing_term_ends_at=now + timedelta(hours=6), + ) + organisation.subscription.subscription_id = "fancy_sub_id23" + organisation.subscription.plan = "startup-v2" + organisation.subscription.save() + OrganisationAPIUsageNotification.objects.create( + organisation=organisation, + percent_usage=100, + notified_at=now, + ) + + mocker.patch("organisations.chargebee.chargebee.chargebee.Subscription.retrieve") + mock_chargebee_update = mocker.patch( + "organisations.chargebee.chargebee.chargebee.Subscription.update" + ) + + mock_api_usage = mocker.patch( + "organisations.tasks.get_current_api_usage", + ) + + mock_api_usage.return_value = 12_005 + assert OrganisationAPIBilling.objects.count() == 0 + + # When + charge_for_api_call_count_overages() + + # Then + mock_chargebee_update.assert_not_called() + assert OrganisationAPIBilling.objects.count() == 0 + + +@pytest.mark.freeze_time("2023-01-19T09:09:47.325132+00:00") +def test_charge_for_api_call_count_overages_with_bad_plan( + organisation: Organisation, + mocker: MockerFixture, +) -> None: + # Given + now = timezone.now() + OrganisationSubscriptionInformationCache.objects.create( + organisation=organisation, + allowed_seats=10, + allowed_projects=3, + allowed_30d_api_calls=10_000, + chargebee_email="test@example.com", + current_billing_term_starts_at=now - timedelta(days=30), + current_billing_term_ends_at=now + timedelta(hours=6), + ) + organisation.subscription.subscription_id = "fancy_sub_id23" + organisation.subscription.plan = "some-bad-plan-someone-randomly-made" + organisation.subscription.save() + OrganisationAPIUsageNotification.objects.create( + organisation=organisation, + percent_usage=100, + notified_at=now, + ) + + mocker.patch("organisations.chargebee.chargebee.chargebee.Subscription.retrieve") + mock_chargebee_update = mocker.patch( + "organisations.chargebee.chargebee.chargebee.Subscription.update" + ) + + mock_api_usage = mocker.patch( + "organisations.tasks.get_current_api_usage", + ) + + mock_api_usage.return_value = 12_005 + assert OrganisationAPIBilling.objects.count() == 0 + + # When + charge_for_api_call_count_overages() + + # Then + # Since the plan is not known ahead of time, it isn't charged. + mock_chargebee_update.assert_not_called() + assert OrganisationAPIBilling.objects.count() == 0 From 89188edd166de4d5dfae05b0002a070d9408fc91 Mon Sep 17 00:00:00 2001 From: Zach Aysan Date: Mon, 8 Apr 2024 18:02:32 +0000 Subject: [PATCH 06/22] Create exceptions for API usage billing --- api/organisations/subscriptions/exceptions.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/api/organisations/subscriptions/exceptions.py b/api/organisations/subscriptions/exceptions.py index 1b24e20bd007..74acd3c6fa92 100644 --- a/api/organisations/subscriptions/exceptions.py +++ b/api/organisations/subscriptions/exceptions.py @@ -22,6 +22,18 @@ class UpgradeSeatsPaymentFailure(APIException): ) +class UpgradeAPIUsageError(APIException): + default_detail = "Failed to upgrade API use in Chargebee" + + +class UpgradeAPIUsagePaymentFailure(APIException): + status_code = 400 + default_detail = ( + "API usage upgrade has failed due to a payment issue. " + "If this persists, contact the organisation admin." + ) + + class SubscriptionDoesNotSupportSeatUpgrade(APIException): status_code = 400 default_detail = "Please Upgrade your plan to add additional seats/users" From 09a4bf04f340741ca89913acdaf1409bb9b965a5 Mon Sep 17 00:00:00 2001 From: Zach Aysan Date: Mon, 8 Apr 2024 18:03:20 +0000 Subject: [PATCH 07/22] Add missing plan constants --- api/organisations/subscriptions/constants.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/api/organisations/subscriptions/constants.py b/api/organisations/subscriptions/constants.py index 690120ec262c..cd3b8ac5ef23 100644 --- a/api/organisations/subscriptions/constants.py +++ b/api/organisations/subscriptions/constants.py @@ -40,6 +40,13 @@ ) FREE_PLAN_ID = "free" TRIAL_SUBSCRIPTION_ID = "trial" +SCALE_UP = "scale-up" +SCALE_UP_12_MONTHS_V2 = "scale-up-12-months-v2" +SCALE_UP_QUARTERLY_V2_SEMIANNUAL = "scale-up-quarterly-v2-semiannual" +SCALE_UP_V2 = "scale-up-v2" +STARTUP = "startup" +STARTUP_ANNUAL_V2 = "startup-annual-v2" +STARTUP_V2 = "startup-v2" class SubscriptionCacheEntity(Enum): From b27c77392d068767beb141e2dcaa1505c0837051 Mon Sep 17 00:00:00 2001 From: Zach Aysan Date: Mon, 8 Apr 2024 18:03:40 +0000 Subject: [PATCH 08/22] Add new addon names --- api/organisations/chargebee/constants.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/api/organisations/chargebee/constants.py b/api/organisations/chargebee/constants.py index e7682d81ebac..f4dbf841bdd9 100644 --- a/api/organisations/chargebee/constants.py +++ b/api/organisations/chargebee/constants.py @@ -1 +1,4 @@ ADDITIONAL_SEAT_ADDON_ID = "additional-team-members-scale-up-v2-monthly" + +ADDITIONAL_API_START_UP_ADDON_ID = "additional-api-start-up-monthly" +ADDITIONAL_API_SCALE_UP_ADDON_ID = "additional-api-scale-up-monthly" From db81e9fca78ce6dc8de807991ef6ae1b6e9a3ab2 Mon Sep 17 00:00:00 2001 From: Zach Aysan Date: Mon, 8 Apr 2024 18:05:45 +0000 Subject: [PATCH 09/22] Add API usage tracking --- api/organisations/chargebee/chargebee.py | 56 +++++++++++++++++++++++- 1 file changed, 55 insertions(+), 1 deletion(-) diff --git a/api/organisations/chargebee/chargebee.py b/api/organisations/chargebee/chargebee.py index faaf9b61a880..8776fdeb3f91 100644 --- a/api/organisations/chargebee/chargebee.py +++ b/api/organisations/chargebee/chargebee.py @@ -11,11 +11,17 @@ from ..subscriptions.constants import CHARGEBEE from ..subscriptions.exceptions import ( CannotCancelChargebeeSubscription, + UpgradeAPIUsageError, + UpgradeAPIUsagePaymentFailure, UpgradeSeatsError, UpgradeSeatsPaymentFailure, ) from .cache import ChargebeeCache -from .constants import ADDITIONAL_SEAT_ADDON_ID +from .constants import ( + ADDITIONAL_API_SCALE_UP_ADDON_ID, + ADDITIONAL_API_START_UP_ADDON_ID, + ADDITIONAL_SEAT_ADDON_ID, +) from .metadata import ChargebeeObjMetadata chargebee.configure(settings.CHARGEBEE_API_KEY, settings.CHARGEBEE_SITE) @@ -203,6 +209,54 @@ def add_single_seat(subscription_id: str): raise UpgradeSeatsError(msg) from e +def add_1000_api_calls_start_up( + subscription_id: str, count: int = 1, invoice_immediately: bool = False +) -> None: + add_1000_api_calls(ADDITIONAL_API_START_UP_ADDON_ID, subscription_id, count) + + +def add_1000_api_calls_scale_up( + subscription_id: str, count: int = 1, invoice_immediately: bool = False +) -> None: + add_1000_api_calls(ADDITIONAL_API_SCALE_UP_ADDON_ID, subscription_id, count) + + +def add_1000_api_calls( + addon_id: str, + subscription_id: str, + count: int = 1, + invoice_immediately: bool = False, +) -> None: + if not count: + return + try: + chargebee.Subscription.update( + subscription_id, + { + "addons": [{"id": addon_id, "quantity": count}], + "prorate": True, + "invoice_immediately": invoice_immediately, + }, + ) + + except ChargebeeAPIError as e: + api_error_code = e.json_obj["api_error_code"] + if api_error_code in CHARGEBEE_PAYMENT_ERROR_CODES: + logger.warning( + f"Payment declined ({api_error_code}) during additional " + f"api calls upgrade to a CB subscription for subscription_id " + f"{subscription_id}" + ) + raise UpgradeAPIUsagePaymentFailure() from e + + msg = ( + "Failed to add additional API calls to CB subscription for subscription id: %s" + % subscription_id + ) + logger.error(msg) + raise UpgradeAPIUsageError(msg) from e + + def _convert_chargebee_subscription_to_dictionary( chargebee_subscription: chargebee.Subscription, ) -> dict: From a7a45c3a0bbe93410d2a89d2a8c15a63668581bf Mon Sep 17 00:00:00 2001 From: Zach Aysan Date: Thu, 18 Apr 2024 15:18:19 +0000 Subject: [PATCH 10/22] Do not prorate --- api/organisations/chargebee/chargebee.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/api/organisations/chargebee/chargebee.py b/api/organisations/chargebee/chargebee.py index 8776fdeb3f91..1f10c9f6630c 100644 --- a/api/organisations/chargebee/chargebee.py +++ b/api/organisations/chargebee/chargebee.py @@ -234,7 +234,7 @@ def add_1000_api_calls( subscription_id, { "addons": [{"id": addon_id, "quantity": count}], - "prorate": True, + "prorate": False, "invoice_immediately": invoice_immediately, }, ) From d469194008590ab5a4bcc25d04b46356b79390c1 Mon Sep 17 00:00:00 2001 From: Zach Aysan Date: Thu, 18 Apr 2024 15:18:38 +0000 Subject: [PATCH 11/22] Update tests to match prorate setting --- api/tests/unit/organisations/test_unit_organisations_tasks.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/api/tests/unit/organisations/test_unit_organisations_tasks.py b/api/tests/unit/organisations/test_unit_organisations_tasks.py index 11ca3c546e97..e28eab4f0761 100644 --- a/api/tests/unit/organisations/test_unit_organisations_tasks.py +++ b/api/tests/unit/organisations/test_unit_organisations_tasks.py @@ -505,7 +505,7 @@ def test_charge_for_api_call_count_overages_scale_up( "quantity": 2, # Two thousand API requests. } ], - "prorate": True, + "prorate": False, "invoice_immediately": False, }, ) @@ -567,7 +567,7 @@ def test_charge_for_api_call_count_overages_start_up( "quantity": 2, # Two thousand API requests. } ], - "prorate": True, + "prorate": False, "invoice_immediately": False, }, ) From e9b58060837122e194392ed49d88d71cffa96a24 Mon Sep 17 00:00:00 2001 From: Zach Aysan Date: Thu, 18 Apr 2024 15:21:47 +0000 Subject: [PATCH 12/22] Add docstring, comments, and an additional select related param --- api/organisations/models.py | 14 ++++++++++++++ api/organisations/tasks.py | 12 ++++++++++++ 2 files changed, 26 insertions(+) diff --git a/api/organisations/models.py b/api/organisations/models.py index f306a825fddb..33d068abd3d2 100644 --- a/api/organisations/models.py +++ b/api/organisations/models.py @@ -481,6 +481,20 @@ class HubspotOrganisation(models.Model): class OrganisationAPIBilling(models.Model): + """ + Tracks API billing for when accounts go over their API usage + limits. This model is what allows subsequent billing runs + to not double bill an organisation for the same use. + + Even though api_overage is charge per thousand API calls, this + class tracks the actual rounded count of API calls that are + billed for (i.e., 52000 for an account with 52233 api calls). + + The option to set immediate_invoice means whether or not the + API billing was processed immediately versus pushed onto the + subsequent subscription billing period. + """ + organisation = models.ForeignKey( Organisation, on_delete=models.CASCADE, related_name="api_billing" ) diff --git a/api/organisations/tasks.py b/api/organisations/tasks.py index 0021daa2cd6a..1984f853190d 100644 --- a/api/organisations/tasks.py +++ b/api/organisations/tasks.py @@ -203,10 +203,21 @@ def handle_api_usage_notifications(): def charge_for_api_call_count_overages(): now = timezone.now() + + # Get the period where we're interested in any new API usage + # notifications for the relevant billing period (ie, this month). api_usage_notified_at = now - timedelta(days=30) + + # Since we're interested in monthly billed accounts, set a wide + # threshold to catch as many billing periods that could be roughly + # considered to be a "monthly" subscription. month_window_start = timedelta(days=25) month_window_end = timedelta(days=35) + + # Only apply charges to ongoing subscriptions that are close to + # being charged due to being at the end of the billing term. closing_billing_term = now + timedelta(hours=12) + organisation_ids = set( OrganisationAPIUsageNotification.objects.filter( notified_at__gte=api_usage_notified_at, @@ -232,6 +243,7 @@ def charge_for_api_call_count_overages(): - month_window_end, ).select_related( "subscription_information_cache", + "subscription", ): subscription_cache = organisation.subscription_information_cache api_usage = get_current_api_usage(organisation.id, "30d") From 50d65abaf53d7462e85f7a1abacf7493cc32a729 Mon Sep 17 00:00:00 2001 From: Zach Aysan Date: Fri, 26 Apr 2024 14:06:07 +0000 Subject: [PATCH 13/22] Tweak comment language --- api/organisations/tasks.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/api/organisations/tasks.py b/api/organisations/tasks.py index 1984f853190d..b4223fae7a67 100644 --- a/api/organisations/tasks.py +++ b/api/organisations/tasks.py @@ -208,9 +208,10 @@ def charge_for_api_call_count_overages(): # notifications for the relevant billing period (ie, this month). api_usage_notified_at = now - timedelta(days=30) - # Since we're interested in monthly billed accounts, set a wide + # Since we're only interested in monthly billed accounts, set a wide # threshold to catch as many billing periods that could be roughly - # considered to be a "monthly" subscription. + # considered to be a "monthly" subscription, while still ruling out + # non-monthly subscriptions. month_window_start = timedelta(days=25) month_window_end = timedelta(days=35) From 5142c39c6589002c28540299c1e4c03ad0ca3718 Mon Sep 17 00:00:00 2001 From: Zach Aysan Date: Fri, 3 May 2024 14:58:03 +0000 Subject: [PATCH 14/22] Add docstring comment as suggested by Matt --- api/organisations/models.py | 1 + 1 file changed, 1 insertion(+) diff --git a/api/organisations/models.py b/api/organisations/models.py index 33d068abd3d2..dfcb92341f1e 100644 --- a/api/organisations/models.py +++ b/api/organisations/models.py @@ -489,6 +489,7 @@ class OrganisationAPIBilling(models.Model): Even though api_overage is charge per thousand API calls, this class tracks the actual rounded count of API calls that are billed for (i.e., 52000 for an account with 52233 api calls). + We're intentionally rounding down to the closest thousands. The option to set immediate_invoice means whether or not the API billing was processed immediately versus pushed onto the From 33251c191c126a977feae983b4db674782f2661a Mon Sep 17 00:00:00 2001 From: Zach Aysan Date: Fri, 3 May 2024 14:59:15 +0000 Subject: [PATCH 15/22] Run the task more frequently --- api/organisations/tasks.py | 5 +++-- .../unit/organisations/test_unit_organisations_tasks.py | 4 ++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/api/organisations/tasks.py b/api/organisations/tasks.py index b4223fae7a67..a49557231ec8 100644 --- a/api/organisations/tasks.py +++ b/api/organisations/tasks.py @@ -217,7 +217,7 @@ def charge_for_api_call_count_overages(): # Only apply charges to ongoing subscriptions that are close to # being charged due to being at the end of the billing term. - closing_billing_term = now + timedelta(hours=12) + closing_billing_term = now + timedelta(hours=1) organisation_ids = set( OrganisationAPIUsageNotification.objects.filter( @@ -284,6 +284,7 @@ def charge_for_api_call_count_overages(): register_recurring_task( run_every=timedelta(hours=12), )(handle_api_usage_notifications) + register_recurring_task( - run_every=timedelta(hours=12), + run_every=timedelta(minutes=30), )(charge_for_api_call_count_overages) diff --git a/api/tests/unit/organisations/test_unit_organisations_tasks.py b/api/tests/unit/organisations/test_unit_organisations_tasks.py index e28eab4f0761..4549eb709e74 100644 --- a/api/tests/unit/organisations/test_unit_organisations_tasks.py +++ b/api/tests/unit/organisations/test_unit_organisations_tasks.py @@ -470,7 +470,7 @@ def test_charge_for_api_call_count_overages_scale_up( allowed_30d_api_calls=10_000, chargebee_email="test@example.com", current_billing_term_starts_at=now - timedelta(days=30), - current_billing_term_ends_at=now + timedelta(hours=6), + current_billing_term_ends_at=now + timedelta(minutes=30), ) organisation.subscription.subscription_id = "fancy_sub_id23" organisation.subscription.plan = "scale-up-v2" @@ -532,7 +532,7 @@ def test_charge_for_api_call_count_overages_start_up( allowed_30d_api_calls=10_000, chargebee_email="test@example.com", current_billing_term_starts_at=now - timedelta(days=30), - current_billing_term_ends_at=now + timedelta(hours=6), + current_billing_term_ends_at=now + timedelta(minutes=30), ) organisation.subscription.subscription_id = "fancy_sub_id23" organisation.subscription.plan = "startup-v2" From 9ed8615ae008b9b0bd14565e2015b7d4a9804c40 Mon Sep 17 00:00:00 2001 From: Zach Aysan Date: Mon, 13 May 2024 13:35:55 +0000 Subject: [PATCH 16/22] Add api calls test for empty given count --- api/organisations/chargebee/__init__.py | 1 + .../test_unit_chargebee_chargebee.py | 24 ++++++++++++++++++- 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/api/organisations/chargebee/__init__.py b/api/organisations/chargebee/__init__.py index cf03c4f2a052..a12b60779abb 100644 --- a/api/organisations/chargebee/__init__.py +++ b/api/organisations/chargebee/__init__.py @@ -1,4 +1,5 @@ from .chargebee import ( # noqa + add_1000_api_calls, add_1000_api_calls_scale_up, add_1000_api_calls_start_up, add_single_seat, diff --git a/api/tests/unit/organisations/chargebee/test_unit_chargebee_chargebee.py b/api/tests/unit/organisations/chargebee/test_unit_chargebee_chargebee.py index d5cea88097f6..1ef2c45c7502 100644 --- a/api/tests/unit/organisations/chargebee/test_unit_chargebee_chargebee.py +++ b/api/tests/unit/organisations/chargebee/test_unit_chargebee_chargebee.py @@ -4,9 +4,11 @@ import pytest from _pytest.monkeypatch import MonkeyPatch from chargebee import APIError +from pytest_mock import MockerFixture from pytz import UTC from organisations.chargebee import ( + add_1000_api_calls, add_single_seat, extract_subscription_metadata, get_customer_id_from_subscription_id, @@ -19,7 +21,10 @@ get_subscription_metadata_from_id, ) from organisations.chargebee.chargebee import cancel_subscription -from organisations.chargebee.constants import ADDITIONAL_SEAT_ADDON_ID +from organisations.chargebee.constants import ( + ADDITIONAL_API_SCALE_UP_ADDON_ID, + ADDITIONAL_SEAT_ADDON_ID, +) from organisations.chargebee.metadata import ChargebeeObjMetadata from organisations.subscriptions.exceptions import ( CannotCancelChargebeeSubscription, @@ -595,3 +600,20 @@ def test_add_single_seat_throws_upgrade_seats_error_error_if_api_error( == "Failed to add additional seat to CB subscription for subscription id: %s" % subscription_id ) + + +def test_add_1000_api_calls_when_count_is_empty(mocker: MockerFixture) -> None: + # Given + subscription_mock = mocker.patch("chargebee.Subscription.update") + + # When + result = add_1000_api_calls( + addon_id=ADDITIONAL_API_SCALE_UP_ADDON_ID, + subscription_id="subscription23", + count=0, + invoice_immediately=True, + ) + + # Then + assert result is None + subscription_mock.assert_not_called() From c6f2b46918038b580e88dd341a86d522a10b4f55 Mon Sep 17 00:00:00 2001 From: Zach Aysan Date: Mon, 13 May 2024 15:39:19 +0000 Subject: [PATCH 17/22] Add organisations tasks test coverage --- .../test_unit_organisations_tasks.py | 112 +++++++++++++++++- 1 file changed, 111 insertions(+), 1 deletion(-) diff --git a/api/tests/unit/organisations/test_unit_organisations_tasks.py b/api/tests/unit/organisations/test_unit_organisations_tasks.py index 6fb47da2e68a..5091fc2a97de 100644 --- a/api/tests/unit/organisations/test_unit_organisations_tasks.py +++ b/api/tests/unit/organisations/test_unit_organisations_tasks.py @@ -1,11 +1,12 @@ import uuid from datetime import timedelta -from unittest.mock import MagicMock +from unittest.mock import MagicMock, call import pytest from django.core.mail.message import EmailMultiAlternatives from django.utils import timezone from freezegun.api import FrozenDateTimeFactory +from pytest_django.fixtures import SettingsWrapper from pytest_mock import MockerFixture from organisations.chargebee.metadata import ChargebeeObjMetadata @@ -32,6 +33,7 @@ charge_for_api_call_count_overages, finish_subscription_cancellation, handle_api_usage_notifications, + register_recurring_tasks, restrict_use_due_to_api_limit_grace_period_over, send_org_over_limit_alert, send_org_subscription_cancelled_alert, @@ -525,6 +527,97 @@ def test_charge_for_api_call_count_overages_scale_up( assert api_billing.billed_at == now +@pytest.mark.freeze_time("2023-01-19T09:09:47.325132+00:00") +def test_charge_for_api_call_count_overages_with_not_covered_plan( + organisation: Organisation, + mocker: MockerFixture, +) -> None: + # Given + now = timezone.now() + OrganisationSubscriptionInformationCache.objects.create( + organisation=organisation, + allowed_seats=10, + allowed_projects=3, + allowed_30d_api_calls=10_000, + chargebee_email="test@example.com", + current_billing_term_starts_at=now - timedelta(days=30), + current_billing_term_ends_at=now + timedelta(minutes=30), + ) + organisation.subscription.subscription_id = "fancy_sub_id23" + + # This plan name is what this test hinges on. + organisation.subscription.plan = "some-plan-not-covered-by-usage" + + organisation.subscription.save() + OrganisationAPIUsageNotification.objects.create( + organisation=organisation, + percent_usage=100, + notified_at=now, + ) + + mocker.patch("organisations.chargebee.chargebee.chargebee.Subscription.retrieve") + mock_chargebee_update = mocker.patch( + "organisations.chargebee.chargebee.chargebee.Subscription.update" + ) + + mock_api_usage = mocker.patch( + "organisations.tasks.get_current_api_usage", + ) + mock_api_usage.return_value = 12_005 + assert OrganisationAPIBilling.objects.count() == 0 + + # When + charge_for_api_call_count_overages() + + # Then + mock_chargebee_update.assert_not_called() + assert OrganisationAPIBilling.objects.count() == 0 + + +@pytest.mark.freeze_time("2023-01-19T09:09:47.325132+00:00") +def test_charge_for_api_call_count_overages_sub_1_api_usage_ratio( + organisation: Organisation, + mocker: MockerFixture, +) -> None: + # Given + now = timezone.now() + OrganisationSubscriptionInformationCache.objects.create( + organisation=organisation, + allowed_seats=10, + allowed_projects=3, + allowed_30d_api_calls=10_000, + chargebee_email="test@example.com", + current_billing_term_starts_at=now - timedelta(days=30), + current_billing_term_ends_at=now + timedelta(minutes=30), + ) + organisation.subscription.subscription_id = "fancy_sub_id23" + organisation.subscription.plan = "scale-up-v2" + organisation.subscription.save() + OrganisationAPIUsageNotification.objects.create( + organisation=organisation, + percent_usage=100, + notified_at=now, + ) + + mocker.patch("organisations.chargebee.chargebee.chargebee.Subscription.retrieve") + mock_chargebee_update = mocker.patch( + "organisations.chargebee.chargebee.chargebee.Subscription.update" + ) + + mock_api_usage = mocker.patch( + "organisations.tasks.get_current_api_usage", + ) + mock_api_usage.return_value = 2_000 + assert OrganisationAPIBilling.objects.count() == 0 + + # When + charge_for_api_call_count_overages() + + # Then + mock_chargebee_update.assert_not_called() + assert OrganisationAPIBilling.objects.count() == 0 + + @pytest.mark.freeze_time("2023-01-19T09:09:47.325132+00:00") def test_charge_for_api_call_count_overages_start_up( organisation: Organisation, @@ -875,3 +968,20 @@ def test_unrestrict_after_api_limit_grace_period_is_stale( assert organisation4.stop_serving_flags is True assert organisation4.block_access_to_admin is True assert getattr(organisation4, "api_limit_access_block", None) is None + + +def test_recurring_tasks(mocker: MockerFixture, settings: SettingsWrapper) -> None: + # Given + settings.ENABLE_API_USAGE_ALERTING = True + register_task_mock = mocker.patch("organisations.tasks.register_recurring_task") + + # When + register_recurring_tasks() + + # Then + register_task_mock.call_args_list == [ + call(run_every=timedelta(seconds=43200)), + call(run_every=timedelta(seconds=1800)), + call(run_every=timedelta(seconds=43200)), + call(run_every=timedelta(seconds=43200)), + ] From b007844f6e56e0bcc090a3cf7b3e0c9ab39e3bd2 Mon Sep 17 00:00:00 2001 From: Zach Aysan Date: Mon, 13 May 2024 15:41:26 +0000 Subject: [PATCH 18/22] Add add_1000_api_calls test coverage --- .../test_unit_chargebee_chargebee.py | 57 +++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/api/tests/unit/organisations/chargebee/test_unit_chargebee_chargebee.py b/api/tests/unit/organisations/chargebee/test_unit_chargebee_chargebee.py index 1ef2c45c7502..e3dae9715336 100644 --- a/api/tests/unit/organisations/chargebee/test_unit_chargebee_chargebee.py +++ b/api/tests/unit/organisations/chargebee/test_unit_chargebee_chargebee.py @@ -4,6 +4,7 @@ import pytest from _pytest.monkeypatch import MonkeyPatch from chargebee import APIError +from chargebee.api_error import APIError as ChargebeeAPIError from pytest_mock import MockerFixture from pytz import UTC @@ -28,6 +29,8 @@ from organisations.chargebee.metadata import ChargebeeObjMetadata from organisations.subscriptions.exceptions import ( CannotCancelChargebeeSubscription, + UpgradeAPIUsageError, + UpgradeAPIUsagePaymentFailure, UpgradeSeatsError, ) @@ -617,3 +620,57 @@ def test_add_1000_api_calls_when_count_is_empty(mocker: MockerFixture) -> None: # Then assert result is None subscription_mock.assert_not_called() + + +def test_add_1000_api_calls_when_chargebee_api_error_has_error_code( + mocker: MockerFixture, +) -> None: + # Given + chargebee_mock = mocker.patch("organisations.chargebee.chargebee.chargebee") + chargebee_response_data = { + "message": "Subscription cannot be created as the payment collection failed. Gateway Error: Card declined.", + "type": "payment", + "api_error_code": "payment_processing_failed", + "param": "item_id", + "error_code": "DeprecatedField", + } + + chargebee_mock.Subscription.update.side_effect = ChargebeeAPIError( + http_code=400, json_obj=chargebee_response_data + ) + + # When / Then + with pytest.raises(UpgradeAPIUsagePaymentFailure): + add_1000_api_calls( + addon_id=ADDITIONAL_API_SCALE_UP_ADDON_ID, + subscription_id="subscription23", + count=1, + invoice_immediately=True, + ) + + +def test_add_1000_api_calls_when_chargebee_api_error_has_no_error_code( + mocker: MockerFixture, +) -> None: + # Given + chargebee_mock = mocker.patch("organisations.chargebee.chargebee.chargebee") + chargebee_response_data = { + "message": "Some massive data failure", + "api_error_code": "halt_and_catch_fire", + "type": "failure", + "param": "item_id", + "error_code": "DeprecatedField", + } + + chargebee_mock.Subscription.update.side_effect = ChargebeeAPIError( + http_code=400, json_obj=chargebee_response_data + ) + + # When / Then + with pytest.raises(UpgradeAPIUsageError): + add_1000_api_calls( + addon_id=ADDITIONAL_API_SCALE_UP_ADDON_ID, + subscription_id="subscription23", + count=1, + invoice_immediately=True, + ) From 0ae1aebb790aa8c03c687f3c2851c8feb9900ddb Mon Sep 17 00:00:00 2001 From: Zach Aysan Date: Mon, 13 May 2024 15:42:28 +0000 Subject: [PATCH 19/22] Add work around for test coverage --- api/organisations/tasks.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/api/organisations/tasks.py b/api/organisations/tasks.py index 37016eec4c45..b9bae7ce1961 100644 --- a/api/organisations/tasks.py +++ b/api/organisations/tasks.py @@ -386,7 +386,12 @@ def unrestrict_after_api_limit_grace_period_is_stale() -> None: organisation.api_limit_access_block.delete() -if settings.ENABLE_API_USAGE_ALERTING: +def register_recurring_tasks() -> None: + """ + Helper function to get codecov coverage. + """ + assert settings.ENABLE_API_USAGE_ALERTING + register_recurring_task( run_every=timedelta(hours=12), )(handle_api_usage_notifications) @@ -402,3 +407,7 @@ def unrestrict_after_api_limit_grace_period_is_stale() -> None: register_recurring_task( run_every=timedelta(hours=12), )(unrestrict_after_api_limit_grace_period_is_stale) + + +if settings.ENABLE_API_USAGE_ALERTING: + register_recurring_tasks() # pragma: no cover From 2c5a43d97f62122b24e233d08ee0298765521367 Mon Sep 17 00:00:00 2001 From: Zach Aysan Date: Mon, 13 May 2024 17:59:50 +0000 Subject: [PATCH 20/22] Trigger build From 23ed02bc87b8410293e13e6c22e0679a2d7e418b Mon Sep 17 00:00:00 2001 From: Zach Aysan Date: Fri, 17 May 2024 09:27:17 -0400 Subject: [PATCH 21/22] Update api/tests/unit/organisations/test_unit_organisations_tasks.py Co-authored-by: Matthew Elwell --- api/tests/unit/organisations/test_unit_organisations_tasks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/api/tests/unit/organisations/test_unit_organisations_tasks.py b/api/tests/unit/organisations/test_unit_organisations_tasks.py index 5091fc2a97de..5434ef629ca2 100644 --- a/api/tests/unit/organisations/test_unit_organisations_tasks.py +++ b/api/tests/unit/organisations/test_unit_organisations_tasks.py @@ -970,7 +970,7 @@ def test_unrestrict_after_api_limit_grace_period_is_stale( assert getattr(organisation4, "api_limit_access_block", None) is None -def test_recurring_tasks(mocker: MockerFixture, settings: SettingsWrapper) -> None: +def test_register_recurring_tasks(mocker: MockerFixture, settings: SettingsWrapper) -> None: # Given settings.ENABLE_API_USAGE_ALERTING = True register_task_mock = mocker.patch("organisations.tasks.register_recurring_task") From f4edbd78c518ef752c6485d0358d16dac6d2fd5f Mon Sep 17 00:00:00 2001 From: Zach Aysan Date: Fri, 17 May 2024 13:42:09 +0000 Subject: [PATCH 22/22] Add assertion for tasks being passed in --- .../organisations/test_unit_organisations_tasks.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/api/tests/unit/organisations/test_unit_organisations_tasks.py b/api/tests/unit/organisations/test_unit_organisations_tasks.py index 5434ef629ca2..8082698ef5c7 100644 --- a/api/tests/unit/organisations/test_unit_organisations_tasks.py +++ b/api/tests/unit/organisations/test_unit_organisations_tasks.py @@ -970,7 +970,9 @@ def test_unrestrict_after_api_limit_grace_period_is_stale( assert getattr(organisation4, "api_limit_access_block", None) is None -def test_register_recurring_tasks(mocker: MockerFixture, settings: SettingsWrapper) -> None: +def test_register_recurring_tasks( + mocker: MockerFixture, settings: SettingsWrapper +) -> None: # Given settings.ENABLE_API_USAGE_ALERTING = True register_task_mock = mocker.patch("organisations.tasks.register_recurring_task") @@ -979,9 +981,18 @@ def test_register_recurring_tasks(mocker: MockerFixture, settings: SettingsWrapp register_recurring_tasks() # Then + # Check when the tasks have been registered register_task_mock.call_args_list == [ call(run_every=timedelta(seconds=43200)), call(run_every=timedelta(seconds=1800)), call(run_every=timedelta(seconds=43200)), call(run_every=timedelta(seconds=43200)), ] + + # And check which tasks were passed in + register_task_mock.return_value.call_args_list == [ + call(handle_api_usage_notifications), + call(charge_for_api_call_count_overages), + call(restrict_use_due_to_api_limit_grace_period_over), + call(unrestrict_after_api_limit_grace_period_is_stale), + ]