From 3225c47043f9647a7426b7f05890bde29b681acc Mon Sep 17 00:00:00 2001 From: Zach Aysan Date: Tue, 6 Aug 2024 15:43:17 -0400 Subject: [PATCH] fix: Set grace period to a singular event (#4455) --- ...eate_organisation_breached_grace_period.py | 38 ++++++++++++ api/organisations/models.py | 8 +++ api/organisations/tasks.py | 24 ++++++-- .../test_unit_organisations_tasks.py | 59 +++++++++++++++++++ 4 files changed, 123 insertions(+), 6 deletions(-) create mode 100644 api/organisations/migrations/0056_create_organisation_breached_grace_period.py diff --git a/api/organisations/migrations/0056_create_organisation_breached_grace_period.py b/api/organisations/migrations/0056_create_organisation_breached_grace_period.py new file mode 100644 index 000000000000..cc51bf632745 --- /dev/null +++ b/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", + ), + ), + ], + ), + ] diff --git a/api/organisations/models.py b/api/organisations/models.py index 3b818c85a463..1451cacc30b3 100644 --- a/api/organisations/models.py +++ b/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" diff --git a/api/organisations/tasks.py b/api/organisations/tasks.py index 97eae8072491..0ebdc1f5d6ca 100644 --- a/api/organisations/tasks.py +++ b/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: diff --git a/api/tests/unit/organisations/test_unit_organisations_tasks.py b/api/tests/unit/organisations/test_unit_organisations_tasks.py index 6c853f097c41..62708193a94a 100644 --- a/api/tests/unit/organisations/test_unit_organisations_tasks.py +++ b/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,