Ensure Subscription Update doesn't update the gifted quotas (#2012)

- add a separate OrgQuotasIn where all quota updates are optional - ensure gifted quotas are never updated as part of org update - update tests
webrecorder · Aug 20, 2024 · 8c9a14b · 8c9a14b
1 parent 351e92a
commit 8c9a14b
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 29 deletions.
diff --git a/backend/btrixcloud/models.py b/backend/btrixcloud/models.py
@@ -1115,12 +1115,28 @@ class S3Storage(BaseModel):
 class OrgQuotas(BaseModel):
     """Organization quotas (settable by superadmin)"""
 
-    maxConcurrentCrawls: Optional[int] = 0
-    maxPagesPerCrawl: Optional[int] = 0
-    storageQuota: Optional[int] = 0
-    maxExecMinutesPerMonth: Optional[int] = 0
-    extraExecMinutes: Optional[int] = 0
-    giftedExecMinutes: Optional[int] = 0
+    storageQuota: int = 0
+    maxExecMinutesPerMonth: int = 0
+
+    maxConcurrentCrawls: int = 0
+    maxPagesPerCrawl: int = 0
+
+    extraExecMinutes: int = 0
+    giftedExecMinutes: int = 0
+
+
+# ============================================================================
+class OrgQuotasIn(BaseModel):
+    """Update for existing OrgQuotas"""
+
+    storageQuota: Optional[int] = None
+    maxExecMinutesPerMonth: Optional[int] = None
+
+    maxConcurrentCrawls: Optional[int] = None
+    maxPagesPerCrawl: Optional[int] = None
+
+    extraExecMinutes: Optional[int] = None
+    giftedExecMinutes: Optional[int] = None
 
 
 # ============================================================================
@@ -1176,7 +1192,7 @@ class SubscriptionUpdate(BaseModel):
     planId: str
 
     futureCancelDate: Optional[datetime] = None
-    quotas: Optional[OrgQuotas] = None
+    quotas: Optional[OrgQuotasIn] = None
 
 
 # ============================================================================

diff --git a/backend/btrixcloud/orgs.py b/backend/btrixcloud/orgs.py
@@ -32,6 +32,7 @@
     Organization,
     StorageRef,
     OrgQuotas,
+    OrgQuotasIn,
     OrgQuotaUpdate,
     OrgReadOnlyUpdate,
     OrgReadOnlyOnCancel,
@@ -489,6 +490,8 @@ async def update_subscription_data(
 
         org = Organization.from_dict(org_data)
         if update.quotas:
+            # don't change gifted minutes here
+            update.quotas.giftedExecMinutes = None
             await self.update_quotas(org, update.quotas)
 
         return org
@@ -512,7 +515,7 @@ async def update_custom_storages(self, org: Organization) -> bool:
         res = await self.orgs.find_one_and_update({"_id": org.id}, {"$set": set_dict})
         return res is not None
 
-    async def update_quotas(self, org: Organization, quotas: OrgQuotas) -> None:
+    async def update_quotas(self, org: Organization, quotas: OrgQuotasIn) -> None:
         """update organization quotas"""
 
         previous_extra_mins = (
@@ -1458,7 +1461,7 @@ async def rename_org(
 
     @router.post("/quotas", tags=["organizations"], response_model=UpdatedResponse)
     async def update_quotas(
-        quotas: OrgQuotas,
+        quotas: OrgQuotasIn,
         org: Organization = Depends(org_owner_dep),
         user: User = Depends(user_dep),
     ):

diff --git a/backend/test/test_org_subs.py b/backend/test/test_org_subs.py
@@ -479,10 +479,10 @@ def test_subscription_events_log(admin_auth_headers, non_default_org_id):
             "quotas": {
                 "maxPagesPerCrawl": 50,
                 "storageQuota": 500000,
-                "extraExecMinutes": 0,
-                "giftedExecMinutes": 0,
-                "maxConcurrentCrawls": 0,
-                "maxExecMinutesPerMonth": 0,
+                "extraExecMinutes": None,
+                "giftedExecMinutes": None,
+                "maxConcurrentCrawls": None,
+                "maxExecMinutesPerMonth": None,
             },
         },
         {"subId": "123", "oid": new_subs_oid, "type": "cancel"},
@@ -547,10 +547,10 @@ def test_subscription_events_log_filter_sub_id(admin_auth_headers):
             "quotas": {
                 "maxPagesPerCrawl": 50,
                 "storageQuota": 500000,
-                "extraExecMinutes": 0,
-                "giftedExecMinutes": 0,
-                "maxConcurrentCrawls": 0,
-                "maxExecMinutesPerMonth": 0,
+                "extraExecMinutes": None,
+                "giftedExecMinutes": None,
+                "maxConcurrentCrawls": None,
+                "maxExecMinutesPerMonth": None,
             },
         },
         {"subId": "123", "oid": new_subs_oid, "type": "cancel"},
@@ -608,10 +608,10 @@ def test_subscription_events_log_filter_oid(admin_auth_headers):
             "quotas": {
                 "maxPagesPerCrawl": 50,
                 "storageQuota": 500000,
-                "extraExecMinutes": 0,
-                "giftedExecMinutes": 0,
-                "maxConcurrentCrawls": 0,
-                "maxExecMinutesPerMonth": 0,
+                "extraExecMinutes": None,
+                "giftedExecMinutes": None,
+                "maxConcurrentCrawls": None,
+                "maxExecMinutesPerMonth": None,
             },
         },
         {"subId": "123", "oid": new_subs_oid, "type": "cancel"},
@@ -643,10 +643,10 @@ def test_subscription_events_log_filter_plan_id(admin_auth_headers):
             "quotas": {
                 "maxPagesPerCrawl": 50,
                 "storageQuota": 500000,
-                "extraExecMinutes": 0,
-                "giftedExecMinutes": 0,
-                "maxConcurrentCrawls": 0,
-                "maxExecMinutesPerMonth": 0,
+                "extraExecMinutes": None,
+                "giftedExecMinutes": None,
+                "maxConcurrentCrawls": None,
+                "maxExecMinutesPerMonth": None,
             },
         }
     ]
@@ -694,10 +694,10 @@ def test_subscription_events_log_filter_status(admin_auth_headers):
             "quotas": {
                 "maxPagesPerCrawl": 50,
                 "storageQuota": 500000,
-                "extraExecMinutes": 0,
-                "giftedExecMinutes": 0,
-                "maxConcurrentCrawls": 0,
-                "maxExecMinutesPerMonth": 0,
+                "extraExecMinutes": None,
+                "giftedExecMinutes": None,
+                "maxConcurrentCrawls": None,
+                "maxExecMinutesPerMonth": None,
             },
         },
     ]