Update logic to ignore scale up plans

Flagsmith · Oct 29, 2024 · 43a4c24 · 43a4c24
1 parent 37313c3
commit 43a4c24
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 10 deletions.
diff --git a/api/organisations/migrations/0058_update_audit_and_history_limits_in_sub_cache.py b/api/organisations/migrations/0058_update_audit_and_history_limits_in_sub_cache.py
@@ -12,7 +12,9 @@
 
 def update_limits(apps: Apps, schema_editor: BaseDatabaseSchemaEditor) -> None:
     subscription_model = apps.get_model("organisations", "Subscription")
-    organisation_subscription_information_cache_model = apps.get_model("organisations", "OrganisationSubscriptionInformationCache")
+    organisation_subscription_information_cache_model = apps.get_model(
+        "organisations", "OrganisationSubscriptionInformationCache"
+    )
 
     all_paid_subscriptions = subscription_model.objects.select_related(
         "organisation", "organisation__subscription_information_cache"
@@ -22,17 +24,17 @@ def update_limits(apps: Apps, schema_editor: BaseDatabaseSchemaEditor) -> None:
 
     for subscription in all_paid_subscriptions:
         subscription_family = SubscriptionPlanFamily.get_by_plan_id(subscription.plan)
-        if subscription_family == SubscriptionPlanFamily.START_UP:
+        if subscription_family != SubscriptionPlanFamily.ENTERPRISE:
+            # We only want to update Enterprise plans since:
+            #  1. start up and scale up should only have the defaults
+            #  2. scale up plans are handled differently (using the VERSIONING_RELEASE_DATE setting) which
+            #     is needed to avoid having to create another plan in chargebee.
             continue
 
-        osic = getattr(subscription.organisation, "subscription_information_cache", None)
-        if osic is None:
+        if (osic := getattr(subscription.organisation, "subscription_information_cache", None)) is None:
             continue
 
-        if subscription_family in (SubscriptionPlanFamily.SCALE_UP, SubscriptionPlanFamily.ENTERPRISE):
-            osic.audit_log_visibility_days = None
-            osic.feature_history_visibility_days = None
-
+        osic.audit_log_visibility_days = osic.feature_history_visibility_days = None
         cache_models_to_update.append(osic)
 
     organisation_subscription_information_cache_model.objects.bulk_update(

diff --git a/api/tests/unit/organisations/permissions/test_unit_organisations_migrations.py b/api/tests/unit/organisations/permissions/test_unit_organisations_migrations.py
@@ -240,11 +240,13 @@ def test_update_audit_and_history_limits(migrator: Migrator) -> None:
     assert migrated_start_up_osic.audit_log_visibility_days == 0
     assert migrated_start_up_osic.feature_history_visibility_days == 7
 
+    # Note that scale up plans are not updated as the logic to grandfather old scale up
+    # plans is handled by the `VERSIONING_RELEASE_DATE` setting.
     migrated_scale_up_osic = NewOrganisationSubscriptionInformationCache.objects.get(
         organisation_id=scale_up_organisation.id,
     )
-    assert migrated_scale_up_osic.audit_log_visibility_days is None
-    assert migrated_scale_up_osic.feature_history_visibility_days is None
+    assert migrated_scale_up_osic.audit_log_visibility_days == 0
+    assert migrated_scale_up_osic.feature_history_visibility_days == 7
 
     migrated_enterprise_osic = NewOrganisationSubscriptionInformationCache.objects.get(
         organisation_id=enterprise_organisation.id,