From ca2b13b06c747d026104de7f012d067905e35a88 Mon Sep 17 00:00:00 2001 From: Zach Aysan Date: Wed, 26 Jun 2024 11:50:20 -0400 Subject: [PATCH] feat(api usage): Extra Flagsmith checks for API overage charges (#4251) Co-authored-by: Matthew Elwell --- api/organisations/tasks.py | 37 +++++-- .../test_unit_organisations_tasks.py | 102 ++++++++++++++++++ 2 files changed, 130 insertions(+), 9 deletions(-) diff --git a/api/organisations/tasks.py b/api/organisations/tasks.py index ef60d76acf89..fd8e7f281746 100644 --- a/api/organisations/tasks.py +++ b/api/organisations/tasks.py @@ -207,7 +207,7 @@ def handle_api_usage_notifications() -> None: flagsmith_client = get_client("local", local_eval=True) for organisation in Organisation.objects.all().select_related( - "subscription_information_cache", + "subscription", "subscription_information_cache" ): feature_enabled = flagsmith_client.get_identity_flags( organisation.flagsmith_identifier, @@ -253,6 +253,8 @@ def charge_for_api_call_count_overages(): ).values_list("organisation_id", flat=True) ) + flagsmith_client = get_client("local", local_eval=True) + for organisation in ( Organisation.objects.filter( id__in=organisation_ids, @@ -275,6 +277,16 @@ def charge_for_api_call_count_overages(): "subscription", ) ): + flags = flagsmith_client.get_identity_flags( + organisation.flagsmith_identifier, + traits={ + "organisation_id": organisation.id, + "subscription.plan": organisation.subscription.plan, + }, + ) + if not flags.is_feature_enabled("api_usage_overage_charges"): + continue + subscription_cache = organisation.subscription_information_cache api_usage = get_current_api_usage(organisation.id, "30d") @@ -343,13 +355,17 @@ def restrict_use_due_to_api_limit_grace_period_over() -> None: organisation_ids = [] for result in queryset: organisation_ids.append(result["organisation"]) - organisations = Organisation.objects.filter( - id__in=organisation_ids, - subscription__plan=FREE_PLAN_ID, - api_limit_access_block__isnull=True, - ).exclude( - stop_serving_flags=True, - block_access_to_admin=True, + organisations = ( + Organisation.objects.filter( + id__in=organisation_ids, + subscription__plan=FREE_PLAN_ID, + api_limit_access_block__isnull=True, + ) + .select_related("subscription") + .exclude( + stop_serving_flags=True, + block_access_to_admin=True, + ) ) update_organisations = [] @@ -359,7 +375,10 @@ def restrict_use_due_to_api_limit_grace_period_over() -> None: for organisation in organisations: flags = flagsmith_client.get_identity_flags( organisation.flagsmith_identifier, - traits={"organisation_id": organisation.id}, + traits={ + "organisation_id": organisation.id, + "subscription.plan": organisation.subscription.plan, + }, ) stop_serving = flags.is_feature_enabled("api_limiting_stop_serving_flags") diff --git a/api/tests/unit/organisations/test_unit_organisations_tasks.py b/api/tests/unit/organisations/test_unit_organisations_tasks.py index accd8c2327d2..1bbfc1513439 100644 --- a/api/tests/unit/organisations/test_unit_organisations_tasks.py +++ b/api/tests/unit/organisations/test_unit_organisations_tasks.py @@ -650,6 +650,11 @@ def test_charge_for_api_call_count_overages_scale_up( "organisations.chargebee.chargebee.chargebee.Subscription.update" ) + get_client_mock = mocker.patch("organisations.tasks.get_client") + client_mock = MagicMock() + get_client_mock.return_value = client_mock + client_mock.get_identity_flags.return_value.is_feature_enabled.return_value = True + mock_api_usage = mocker.patch( "organisations.tasks.get_current_api_usage", ) @@ -682,6 +687,63 @@ 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_scale_up_when_flagsmith_client_sets_is_enabled_to_false( + organisation: Organisation, + mocker: MockerFixture, +) -> None: + # Given + now = timezone.now() + OrganisationSubscriptionInformationCache.objects.create( + organisation=organisation, + allowed_seats=10, + allowed_projects=3, + allowed_30d_api_calls=100_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" + ) + + get_client_mock = mocker.patch("organisations.tasks.get_client") + client_mock = MagicMock() + get_client_mock.return_value = client_mock + client_mock.get_identity_flags.return_value.is_feature_enabled.return_value = False + + mock_api_usage = mocker.patch( + "organisations.tasks.get_current_api_usage", + ) + mock_api_usage.return_value = 212_005 + assert OrganisationAPIBilling.objects.count() == 0 + + # When + charge_for_api_call_count_overages() + + # Then + # No charges are applied to the account. + client_mock.get_identity_flags.assert_called_once_with( + organisation.flagsmith_identifier, + traits={ + "organisation_id": organisation.id, + "subscription.plan": organisation.subscription.plan, + }, + ) + 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_grace_period( organisation: Organisation, @@ -707,6 +769,11 @@ def test_charge_for_api_call_count_overages_grace_period( notified_at=now, ) + get_client_mock = mocker.patch("organisations.tasks.get_client") + client_mock = MagicMock() + get_client_mock.return_value = client_mock + client_mock.get_identity_flags.return_value.is_feature_enabled.return_value = True + mock_chargebee_update = mocker.patch( "organisations.chargebee.chargebee.chargebee.Subscription.update" ) @@ -753,6 +820,11 @@ def test_charge_for_api_call_count_overages_with_not_covered_plan( notified_at=now, ) + get_client_mock = mocker.patch("organisations.tasks.get_client") + client_mock = MagicMock() + get_client_mock.return_value = client_mock + client_mock.get_identity_flags.return_value.is_feature_enabled.return_value = True + mocker.patch("organisations.chargebee.chargebee.chargebee.Subscription.retrieve") mock_chargebee_update = mocker.patch( "organisations.chargebee.chargebee.chargebee.Subscription.update" @@ -797,6 +869,11 @@ def test_charge_for_api_call_count_overages_under_api_limit( notified_at=now, ) + get_client_mock = mocker.patch("organisations.tasks.get_client") + client_mock = MagicMock() + get_client_mock.return_value = client_mock + client_mock.get_identity_flags.return_value.is_feature_enabled.return_value = True + mocker.patch("organisations.chargebee.chargebee.chargebee.Subscription.retrieve") mock_chargebee_update = mocker.patch( "organisations.chargebee.chargebee.chargebee.Subscription.update" @@ -841,6 +918,10 @@ def test_charge_for_api_call_count_overages_start_up( notified_at=now, ) + get_client_mock = mocker.patch("organisations.tasks.get_client") + client_mock = MagicMock() + get_client_mock.return_value = client_mock + client_mock.get_identity_flags.return_value.is_feature_enabled.return_value = True mocker.patch("organisations.chargebee.chargebee.chargebee.Subscription.retrieve") mock_chargebee_update = mocker.patch( "organisations.chargebee.chargebee.chargebee.Subscription.update" @@ -918,6 +999,10 @@ def test_charge_for_api_call_count_overages_start_up_with_api_billing( billed_at=now, ) + get_client_mock = mocker.patch("organisations.tasks.get_client") + client_mock = MagicMock() + get_client_mock.return_value = client_mock + client_mock.get_identity_flags.return_value.is_feature_enabled.return_value = True mocker.patch("organisations.chargebee.chargebee.chargebee.Subscription.retrieve") mock_chargebee_update = mocker.patch( "organisations.chargebee.chargebee.chargebee.Subscription.update" @@ -1137,6 +1222,23 @@ def test_restrict_use_due_to_api_limit_grace_period_over( assert organisation2.block_access_to_admin is True assert organisation2.api_limit_access_block + client_mock.get_identity_flags.call_args_list == [ + call( + f"org.{organisation.id}", + traits={ + "organisation_id": organisation.id, + "subscription.plan": organisation.subscription.plan, + }, + ), + call( + f"org.{organisation2.id}", + traits={ + "organisation_id": organisation2.id, + "subscription.plan": organisation2.subscription.plan, + }, + ), + ] + # Organisations that change their subscription are unblocked. organisation.subscription.plan = "scale-up-v2" organisation.subscription.save()