From b2a18995663b3e894a4f8510016ee1139a9685a6 Mon Sep 17 00:00:00 2001
From: Zach Aysan <zachaysan@gmail.com>
Date: Fri, 6 Sep 2024 09:46:46 -0400
Subject: [PATCH] feat: Add subscription cache for new organisations (#4587)

---
 api/organisations/models.py                   |  5 +++++
 .../test_unit_app_analytics_views.py          |  1 +
 .../test_unit_organisations_models.py         | 21 +++++++++++++++++++
 .../test_unit_organisations_tasks.py          |  2 --
 4 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/api/organisations/models.py b/api/organisations/models.py
index 4665ad668742..b87217b2181c 100644
--- a/api/organisations/models.py
+++ b/api/organisations/models.py
@@ -149,6 +149,11 @@ def cancel_subscription(self):
     def create_subscription(self):
         Subscription.objects.create(organisation=self)
 
+    @hook(AFTER_CREATE)
+    def create_subscription_cache(self):
+        if is_saas() and not self.has_subscription_information_cache():
+            OrganisationSubscriptionInformationCache.objects.create(organisation=self)
+
     @hook(AFTER_SAVE)
     def clear_environment_caches(self):
         from environments.models import Environment
diff --git a/api/tests/unit/app_analytics/test_unit_app_analytics_views.py b/api/tests/unit/app_analytics/test_unit_app_analytics_views.py
index 371f54bf329a..8b14e9f3e015 100644
--- a/api/tests/unit/app_analytics/test_unit_app_analytics_views.py
+++ b/api/tests/unit/app_analytics/test_unit_app_analytics_views.py
@@ -231,6 +231,7 @@ def test_get_usage_data__ninety_day_period(
     week_from_now = now + timedelta(days=7)
     four_weeks_ago = now - timedelta(days=28)
     ninety_days_ago = now - timedelta(days=90)
+
     OrganisationSubscriptionInformationCache.objects.create(
         organisation=organisation,
         current_billing_term_starts_at=four_weeks_ago,
diff --git a/api/tests/unit/organisations/test_unit_organisations_models.py b/api/tests/unit/organisations/test_unit_organisations_models.py
index 1d90e690bab5..cc015a85d671 100644
--- a/api/tests/unit/organisations/test_unit_organisations_models.py
+++ b/api/tests/unit/organisations/test_unit_organisations_models.py
@@ -551,3 +551,24 @@ def test_reset_of_api_notifications(organisation: Organisation) -> None:
     # Then
     assert OrganisationAPIUsageNotification.objects.count() == 1
     assert OrganisationAPIUsageNotification.objects.first() == oapiun
+
+
+def test_organisation_creates_subscription_cache(
+    db: None, mocker: MockerFixture
+) -> None:
+    # Given
+    mocker.patch("organisations.models.is_saas", return_value=True)
+
+    # When
+    organisation = Organisation.objects.create(name="Test org")
+
+    # Then
+    assert organisation.subscription_information_cache
+    assert (
+        organisation.subscription_information_cache.allowed_seats
+        == MAX_SEATS_IN_FREE_PLAN
+    )
+    assert (
+        organisation.subscription_information_cache.allowed_30d_api_calls
+        == MAX_API_CALLS_IN_FREE_PLAN
+    )
diff --git a/api/tests/unit/organisations/test_unit_organisations_tasks.py b/api/tests/unit/organisations/test_unit_organisations_tasks.py
index b19b3d9b55d1..a307976008bb 100644
--- a/api/tests/unit/organisations/test_unit_organisations_tasks.py
+++ b/api/tests/unit/organisations/test_unit_organisations_tasks.py
@@ -756,7 +756,6 @@ def test_handle_api_usage_notifications_missing_info_cache(
     from organisations.task_helpers import logger
 
     logger.addHandler(inspecting_handler)
-
     assert organisation.has_subscription_information_cache() is False
 
     mock_api_usage = mocker.patch(
@@ -1264,7 +1263,6 @@ def test_charge_for_api_call_count_overages_with_exception(
     from organisations.tasks import logger
 
     logger.addHandler(inspecting_handler)
-
     OrganisationSubscriptionInformationCache.objects.create(
         organisation=organisation,
         allowed_seats=10,