add OrgQuotasIn with all options optional, fix tests

webrecorder · Aug 10, 2024 · 0e4ada0 · 0e4ada0
1 parent bf26bcd
commit 0e4ada0
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 17 deletions.
diff --git a/backend/btrixcloud/models.py b/backend/btrixcloud/models.py
@@ -1115,10 +1115,25 @@ 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
+    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
@@ -1177,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,
@@ -514,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 = (
@@ -1460,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
@@ -481,8 +481,8 @@ def test_subscription_events_log(admin_auth_headers, non_default_org_id):
                 "storageQuota": 500000,
                 "extraExecMinutes": None,
                 "giftedExecMinutes": None,
-                "maxConcurrentCrawls": 0,
-                "maxExecMinutesPerMonth": 0,
+                "maxConcurrentCrawls": None,
+                "maxExecMinutesPerMonth": None,
             },
         },
         {"subId": "123", "oid": new_subs_oid, "type": "cancel"},
@@ -549,8 +549,8 @@ def test_subscription_events_log_filter_sub_id(admin_auth_headers):
                 "storageQuota": 500000,
                 "extraExecMinutes": None,
                 "giftedExecMinutes": None,
-                "maxConcurrentCrawls": 0,
-                "maxExecMinutesPerMonth": 0,
+                "maxConcurrentCrawls": None,
+                "maxExecMinutesPerMonth": None,
             },
         },
         {"subId": "123", "oid": new_subs_oid, "type": "cancel"},
@@ -610,8 +610,8 @@ def test_subscription_events_log_filter_oid(admin_auth_headers):
                 "storageQuota": 500000,
                 "extraExecMinutes": None,
                 "giftedExecMinutes": None,
-                "maxConcurrentCrawls": 0,
-                "maxExecMinutesPerMonth": 0,
+                "maxConcurrentCrawls": None,
+                "maxExecMinutesPerMonth": None,
             },
         },
         {"subId": "123", "oid": new_subs_oid, "type": "cancel"},
@@ -645,8 +645,8 @@ def test_subscription_events_log_filter_plan_id(admin_auth_headers):
                 "storageQuota": 500000,
                 "extraExecMinutes": None,
                 "giftedExecMinutes": None,
-                "maxConcurrentCrawls": 0,
-                "maxExecMinutesPerMonth": 0,
+                "maxConcurrentCrawls": None,
+                "maxExecMinutesPerMonth": None,
             },
         }
     ]
@@ -696,8 +696,8 @@ def test_subscription_events_log_filter_status(admin_auth_headers):
                 "storageQuota": 500000,
                 "extraExecMinutes": None,
                 "giftedExecMinutes": None,
-                "maxConcurrentCrawls": 0,
-                "maxExecMinutesPerMonth": 0,
+                "maxConcurrentCrawls": None,
+                "maxExecMinutesPerMonth": None,
             },
         },
     ]