Added handling for lifetime quota changes

cellarium/cas/client.py

@@ -232,9 +232,14 @@ def print_user_quota(self) -> None:
         Print the user's quota information
         """
         user_quota = self.cas_api_service.get_user_quota()
+        lifetime_quota = user_quota["lifetime_quota"] or "unlimited"
+        remaining_lifetime_quota = (
+            user_quota["remaining_lifetime_quota"] or "unlimited"
+        )
         self.__print(
-            f"User quota: {user_quota['quota']}, Remaining quota: {user_quota['remaining_quota']}, "
-            f"Reset date: {user_quota['quota_reset_date']}"
+            f"Weekly quota: {user_quota['weekly_quota']}, Remaining weekly quota: {user_quota['remaining_weekly_quota']}, "
+            f"Weekly quota reset date: {user_quota['quota_reset_date']}\n"
+            f"Lifetime quota: {lifetime_quota}, Remaining lifetime quota: {remaining_lifetime_quota} "
         )
 
     def __get_async_sharded_request_callback(
@@ -445,9 +450,14 @@ def __validate_cells_under_quota(self, cell_count: int) -> None:
         :param cell_count: Number of cells in the input data
         """
         user_quota = self.cas_api_service.get_user_quota()
-        if cell_count > user_quota["remaining_quota"]:
+        if cell_count > (user_quota["remaining_lifetime_quota"] or cell_count):
+            raise exceptions.QuotaExceededError(
+                f"Number of cells in the input data ({cell_count}) exceeds the user's remaining lifetime quota ({user_quota['remaining_lifetime_quota']}).  "
+                f"If you would like to discuss removing your lifetime quota, please reach out to [email protected]"
+            )
+        elif cell_count > user_quota["remaining_weekly_quota"]:
             raise exceptions.QuotaExceededError(
-                f"Number of cells in the input data ({cell_count}) exceeds the user's remaining quota ({user_quota['remaining_quota']}).  "
+                f"Number of cells in the input data ({cell_count}) exceeds the user's remaining quota ({user_quota['remaining_weekly_quota']}).  "
                 f"The user's quota will be reset to {user_quota['quota']} on {user_quota['quota_reset_date']}."
             )
 

tests/unit/test_cas_client.py

@@ -444,8 +444,10 @@ def __mock_pre_call_requests(self, num_features: int = 3):
             status_code=200,
             response_body={
                 "user_id": 0,
-                "quota": 1000,
-                "remaining_quota": 1000,
+                "weekly_quota": 1000,
+                "remaining_weekly_quota": 1000,
+                "lifetime_quota": None,
+                "remaining_lifetime_quota": None,
                 "quota_reset_date": datetime.datetime.today() + 7 * datetime.timedelta(days=1),
             },
         )