From d273679b3236c3fceccfeb71c401df5a2d69ad27 Mon Sep 17 00:00:00 2001
From: Matthew Elwell <matthew.elwell@flagsmith.com>
Date: Mon, 19 Aug 2024 18:56:44 +0100
Subject: [PATCH] fix: subscription info cache race condition (#4518)

---
 api/organisations/models.py                   | 26 ++++++++++++------
 .../test_unit_organisations_tasks.py          | 27 ++++++++++++++++++-
 2 files changed, 44 insertions(+), 9 deletions(-)

diff --git a/api/organisations/models.py b/api/organisations/models.py
index 1451cacc30b3..4665ad668742 100644
--- a/api/organisations/models.py
+++ b/api/organisations/models.py
@@ -303,18 +303,14 @@ def save_as_free_subscription(self):
         if not getattr(self.organisation, "subscription_information_cache", None):
             return
 
-        # There is a weird bug where the cache is present, but the id is unset.
-        # See here for more: https://flagsmith.sentry.io/issues/4945988284/
-        if not self.organisation.subscription_information_cache.id:
-            return
-
-        self.organisation.subscription_information_cache.delete()
+        self.organisation.subscription_information_cache.reset_to_defaults()
+        self.organisation.subscription_information_cache.save()
 
     def prepare_for_cancel(
         self, cancellation_date=timezone.now(), update_chargebee=True
     ) -> None:
         """
-        This method get's a subscription ready for cancelation.
+        This method gets a subscription ready for cancellation.
 
         If cancellation_date is in the future some aspects are
         reserved for a task after the date has passed.
@@ -451,7 +447,7 @@ class OrganisationSubscriptionInformationCache(LifecycleModelMixin, models.Model
     api_calls_7d = models.IntegerField(default=0)
     api_calls_30d = models.IntegerField(default=0)
 
-    allowed_seats = models.IntegerField(default=1)
+    allowed_seats = models.IntegerField(default=MAX_SEATS_IN_FREE_PLAN)
     allowed_30d_api_calls = models.IntegerField(default=MAX_API_CALLS_IN_FREE_PLAN)
     allowed_projects = models.IntegerField(default=1, blank=True, null=True)
 
@@ -461,6 +457,20 @@ class OrganisationSubscriptionInformationCache(LifecycleModelMixin, models.Model
     def erase_api_notifications(self):
         self.organisation.api_usage_notifications.all().delete()
 
+    def reset_to_defaults(self):
+        """
+        Resets all limits and CB related data to the defaults, leaving the
+        usage data intact.
+        """
+        self.current_billing_term_starts_at = None
+        self.current_billing_term_ends_at = None
+
+        self.allowed_seats = MAX_SEATS_IN_FREE_PLAN
+        self.allowed_30d_api_calls = MAX_API_CALLS_IN_FREE_PLAN
+        self.allowed_projects = 1
+
+        self.chargebee_email = None
+
 
 class OrganisationAPIUsageNotification(models.Model):
     organisation = models.ForeignKey(
diff --git a/api/tests/unit/organisations/test_unit_organisations_tasks.py b/api/tests/unit/organisations/test_unit_organisations_tasks.py
index 6ab939aca3f1..d97261447a2f 100644
--- a/api/tests/unit/organisations/test_unit_organisations_tasks.py
+++ b/api/tests/unit/organisations/test_unit_organisations_tasks.py
@@ -109,6 +109,13 @@ def test_subscription_cancellation(db: None) -> None:
     organisation = Organisation.objects.create()
     OrganisationSubscriptionInformationCache.objects.create(
         organisation=organisation,
+        allowed_seats=5,
+        allowed_30d_api_calls=1_000_000,
+        allowed_projects=None,
+        api_calls_24h=30_000,
+        api_calls_7d=210_000,
+        api_calls_30d=900_000,
+        chargebee_email="foo@example.com",
     )
     UserOrganisation.objects.create(
         organisation=organisation,
@@ -138,7 +145,8 @@ def test_subscription_cancellation(db: None) -> None:
     # Then
     organisation.refresh_from_db()
     subscription.refresh_from_db()
-    assert getattr(organisation, "subscription_information_cache", None) is None
+    organisation.subscription_information_cache.refresh_from_db()
+
     assert subscription.subscription_id is None
     assert subscription.subscription_date is None
     assert subscription.plan == FREE_PLAN_ID
@@ -151,6 +159,23 @@ def test_subscription_cancellation(db: None) -> None:
     assert subscription.cancellation_date is None
     assert subscription.notes == notes
 
+    # The CB / limit data on the subscription information cache object is reset
+    assert organisation.subscription_information_cache.chargebee_email is None
+    assert (
+        organisation.subscription_information_cache.allowed_30d_api_calls
+        == MAX_API_CALLS_IN_FREE_PLAN
+    )
+    assert organisation.subscription_information_cache.allowed_projects == 1
+    assert (
+        organisation.subscription_information_cache.allowed_seats
+        == MAX_SEATS_IN_FREE_PLAN
+    )
+
+    # But the usage data isn't
+    assert organisation.subscription_information_cache.api_calls_24h == 30_000
+    assert organisation.subscription_information_cache.api_calls_7d == 210_000
+    assert organisation.subscription_information_cache.api_calls_30d == 900_000
+
 
 @pytest.mark.freeze_time("2023-01-19T09:12:34+00:00")
 def test_finish_subscription_cancellation(db: None, mocker: MockerFixture) -> None: