fix: Set grace period to a singular event (#4455)

api/organisations/migrations/0056_create_organisation_breached_grace_period.py

@@ -0,0 +1,38 @@
+# Generated by Django 3.2.25 on 2024-08-06 17:46
+
+import django.db.models.deletion
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("organisations", "0055_alter_percent_usage"),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name="OrganisationBreachedGracePeriod",
+            fields=[
+                (
+                    "id",
+                    models.AutoField(
+                        auto_created=True,
+                        primary_key=True,
+                        serialize=False,
+                        verbose_name="ID",
+                    ),
+                ),
+                ("created_at", models.DateTimeField(auto_now_add=True, null=True)),
+                ("updated_at", models.DateTimeField(auto_now=True, null=True)),
+                (
+                    "organisation",
+                    models.OneToOneField(
+                        on_delete=django.db.models.deletion.CASCADE,
+                        related_name="breached_grace_period",
+                        to="organisations.organisation",
+                    ),
+                ),
+            ],
+        ),
+    ]

api/organisations/models.py

@@ -476,6 +476,14 @@ class OrganisationAPIUsageNotification(models.Model):
     updated_at = models.DateTimeField(null=True, auto_now=True)
 
 
+class OrganisationBreachedGracePeriod(models.Model):
+    organisation = models.OneToOneField(
+        Organisation, on_delete=models.CASCADE, related_name="breached_grace_period"
+    )
+    created_at = models.DateTimeField(null=True, auto_now_add=True)
+    updated_at = models.DateTimeField(null=True, auto_now=True)
+
+
 class APILimitAccessBlock(models.Model):
     organisation = models.OneToOneField(
         Organisation, on_delete=models.CASCADE, related_name="api_limit_access_block"

api/organisations/tasks.py

@@ -4,7 +4,7 @@
 
 from app_analytics.influxdb_wrapper import get_current_api_usage
 from django.conf import settings
-from django.db.models import F, Max
+from django.db.models import F, Max, Q
 from django.utils import timezone
 from task_processor.decorators import (
     register_recurring_task,
@@ -22,6 +22,7 @@
     Organisation,
     OrganisationAPIBilling,
     OrganisationAPIUsageNotification,
+    OrganisationBreachedGracePeriod,
     Subscription,
 )
 from organisations.subscriptions.constants import FREE_PLAN_ID
@@ -243,13 +244,22 @@ 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)
+    now = timezone.now()
+    grace_period = now - timedelta(days=API_USAGE_GRACE_PERIOD)
+    month_start = now - timedelta(30)
     queryset = (
         OrganisationAPIUsageNotification.objects.filter(
-            notified_at__gt=month_start,
-            notified_at__lt=grace_period,
-            percent_usage__gte=100,
+            Q(
+                notified_at__gte=month_start,
+                notified_at__lte=grace_period,
+                percent_usage__gte=100,
+            )
+            | Q(
+                notified_at__gte=month_start,
+                notified_at__lte=now,
+                percent_usage__gte=100,
+                organisation__breached_grace_period__isnull=False,
+            )
         )
         .values("organisation")
         .annotate(max_value=Max("percent_usage"))
@@ -293,6 +303,8 @@ def restrict_use_due_to_api_limit_grace_period_over() -> None:
         if not organisation.has_subscription_information_cache():
             continue
 
+        OrganisationBreachedGracePeriod.objects.get_or_create(organisation=organisation)
+
         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:

api/tests/unit/organisations/test_unit_organisations_tasks.py

@@ -21,6 +21,7 @@
     Organisation,
     OrganisationAPIBilling,
     OrganisationAPIUsageNotification,
+    OrganisationBreachedGracePeriod,
     OrganisationRole,
     OrganisationSubscriptionInformationCache,
     UserOrganisation,
@@ -1410,6 +1411,64 @@ 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_breached(
+    mocker: MockerFixture,
+    organisation: Organisation,
+    freezer: FrozenDateTimeFactory,
+    mailoutbox: list[EmailMultiAlternatives],
+    admin_user: FFAdminUser,
+    staff_user: FFAdminUser,
+) -> None:
+    # Given
+    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()
+
+    OrganisationBreachedGracePeriod.objects.create(organisation=organisation)
+    OrganisationSubscriptionInformationCache.objects.create(
+        organisation=organisation,
+        allowed_seats=10,
+        allowed_projects=3,
+        allowed_30d_api_calls=10_000,
+        chargebee_email="test@example.com",
+    )
+    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 = 12_005
+
+    OrganisationAPIUsageNotification.objects.create(
+        notified_at=now,
+        organisation=organisation,
+        percent_usage=100,
+    )
+    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()
+
+    assert organisation.stop_serving_flags is True
+    assert organisation.block_access_to_admin is True
+    assert organisation.api_limit_access_block
+
+
 @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,