feat(api usage): Extra Flagsmith checks for API overage charges (#4251)

Co-authored-by: Matthew Elwell <[email protected]>

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")

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="[email protected]",
+        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()