fix: Check API usage before restricting serving flags and admin (#4422)

api/organisations/tasks.py

@@ -243,7 +243,6 @@ def restrict_use_due_to_api_limit_grace_period_over() -> None:
     Since free plans don't have predefined subscription periods, we
     use a rolling thirty day period to filter them.
     """
-
     grace_period = timezone.now() - timedelta(days=API_USAGE_GRACE_PERIOD)
     month_start = timezone.now() - timedelta(30)
     queryset = (
@@ -259,13 +258,14 @@ 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,
         )
-        .select_related("subscription")
+        .select_related("subscription", "subscription_information_cache")
         .exclude(
             stop_serving_flags=True,
             block_access_to_admin=True,
@@ -290,6 +290,17 @@ def restrict_use_due_to_api_limit_grace_period_over() -> None:
         if not stop_serving and not block_access:
             continue
 
+        if not organisation.has_subscription_information_cache():
+            continue
+
+        subscription_cache = organisation.subscription_information_cache
+        api_usage = get_current_api_usage(organisation.id, "30d")
+        if api_usage / subscription_cache.allowed_30d_api_calls < 1.0:
+            logger.info(
+                f"API use for organisation {organisation.id} has fallen to below limit, so not restricting use."
+            )
+            continue
+
         organisation.stop_serving_flags = stop_serving
         organisation.block_access_to_admin = block_access
 

api/tests/unit/organisations/test_unit_organisations_tasks.py

@@ -1162,6 +1162,29 @@ def test_restrict_use_due_to_api_limit_grace_period_over(
     organisation4 = Organisation.objects.create(name="Org #4")
     organisation5 = Organisation.objects.create(name="Org #5")
 
+    for org in [
+        organisation,
+        organisation2,
+        organisation3,
+        organisation4,
+        organisation5,
+    ]:
+        OrganisationSubscriptionInformationCache.objects.create(
+            organisation=org,
+            allowed_seats=10,
+            allowed_projects=3,
+            allowed_30d_api_calls=10_000,
+            chargebee_email="[email protected]",
+        )
+        org.subscription.subscription_id = "fancy_sub_id23"
+        org.subscription.plan = FREE_PLAN_ID
+        org.subscription.save()
+
+    mock_api_usage = mocker.patch(
+        "organisations.tasks.get_current_api_usage",
+    )
+    mock_api_usage.return_value = 12_005
+
     # Add users to test email delivery
     for org in [organisation2, organisation3, organisation4, organisation5]:
         admin_user.add_organisation(org, role=OrganisationRole.ADMIN)
@@ -1296,6 +1319,111 @@ def test_restrict_use_due_to_api_limit_grace_period_over(
     assert getattr(organisation, "api_limit_access_block", None) is None
 
 
+@pytest.mark.freeze_time("2023-01-19T09:09:47.325132+00:00")
+def test_restrict_use_due_to_api_limit_grace_period_over_missing_subscription_information_cache(
+    mocker: MockerFixture,
+    organisation: Organisation,
+    freezer: FrozenDateTimeFactory,
+    mailoutbox: list[EmailMultiAlternatives],
+) -> None:
+    # Given
+    assert not organisation.has_subscription_information_cache()
+
+    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
+
+    now = timezone.now()
+    organisation.subscription.subscription_id = "fancy_sub_id23"
+    organisation.subscription.plan = FREE_PLAN_ID
+    organisation.subscription.save()
+
+    OrganisationAPIUsageNotification.objects.create(
+        notified_at=now,
+        organisation=organisation,
+        percent_usage=120,
+    )
+
+    now = now + timedelta(days=API_USAGE_GRACE_PERIOD + 1)
+    freezer.move_to(now)
+
+    # When
+    restrict_use_due_to_api_limit_grace_period_over()
+
+    # Then
+    organisation.refresh_from_db()
+
+    # Organisations that missing the cache don't get blocked
+    assert organisation.stop_serving_flags is False
+    assert organisation.block_access_to_admin is False
+    assert not hasattr(organisation, "api_limit_access_block")
+    assert len(mailoutbox) == 0
+
+
+@pytest.mark.freeze_time("2023-01-19T09:09:47.325132+00:00")
+def test_restrict_use_due_to_api_limit_grace_period_over_with_reduced_api_usage(
+    mocker: MockerFixture,
+    organisation: Organisation,
+    freezer: FrozenDateTimeFactory,
+    mailoutbox: list[EmailMultiAlternatives],
+    inspecting_handler: logging.Handler,
+) -> None:
+    # Given
+    assert not organisation.has_subscription_information_cache()
+
+    from organisations.tasks import logger
+
+    logger.addHandler(inspecting_handler)
+
+    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
+
+    now = timezone.now()
+
+    OrganisationSubscriptionInformationCache.objects.create(
+        organisation=organisation,
+        allowed_seats=10,
+        allowed_projects=3,
+        allowed_30d_api_calls=10_000,
+        chargebee_email="[email protected]",
+    )
+    organisation.subscription.subscription_id = "fancy_sub_id23"
+    organisation.subscription.plan = FREE_PLAN_ID
+    organisation.subscription.save()
+
+    mock_api_usage = mocker.patch(
+        "organisations.tasks.get_current_api_usage",
+    )
+    mock_api_usage.return_value = 8000
+
+    OrganisationAPIUsageNotification.objects.create(
+        notified_at=now,
+        organisation=organisation,
+        percent_usage=120,
+    )
+
+    now = now + timedelta(days=API_USAGE_GRACE_PERIOD + 1)
+    freezer.move_to(now)
+
+    # When
+    restrict_use_due_to_api_limit_grace_period_over()
+
+    # Then
+    organisation.refresh_from_db()
+
+    # Organisations that missing the cache don't get blocked
+    assert organisation.stop_serving_flags is False
+    assert organisation.block_access_to_admin is False
+    assert not hasattr(organisation, "api_limit_access_block")
+    assert len(mailoutbox) == 0
+    assert inspecting_handler.messages == [
+        f"API use for organisation {organisation.id} has fallen to below limit, so not restricting use."
+    ]
+
+
 @pytest.mark.freeze_time("2023-01-19T09:09:47.325132+00:00")
 def test_unrestrict_after_api_limit_grace_period_is_stale(
     organisation: Organisation,