fix deploy tests

dfinity · Feb 6, 2025 · 9b53ff1 · 9b53ff1
1 parent 48706a8
commit 9b53ff1
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 13 deletions.
diff --git a/core/control-panel/impl/src/services/deploy.rs b/core/control-panel/impl/src/services/deploy.rs
@@ -4,8 +4,8 @@ use crate::{
         canister_config, CallContext, CANISTER_CREATION_CYCLES, INITIAL_STATION_CYCLES,
         INITIAL_UPGRADER_CYCLES, NNS_ROOT_CANISTER_ID,
     },
-    errors::DeployError,
-    models::UserStation,
+    errors::{DeployError, UserError},
+    models::{CanDeployStation, UserStation},
     services::{USER_SERVICE, USER_STATION_SERVICE},
 };
 use candid::{Encode, Principal};
@@ -57,7 +57,18 @@ impl DeployService {
         let station_wasm_module = config.station_wasm_module;
         let station_wasm_module_extra_chunks = config.station_wasm_module_extra_chunks;
 
-        self.user_service.can_deploy_station(ctx)?;
+        let can_deploy_station_response = user.can_deploy_station();
+        match can_deploy_station_response {
+            CanDeployStation::Allowed(_) => {}
+            CanDeployStation::QuotaExceeded => {
+                return Err(UserError::DeployStationQuotaExceeded)?;
+            }
+            CanDeployStation::NotAllowed(subscription_status) => {
+                return Err(UserError::BadUserSubscriptionStatus {
+                    subscription_status: subscription_status.into(),
+                })?;
+            }
+        }
 
         // Creates the station canister
         let station_canister = create_canister(input.subnet_selection, CANISTER_CREATION_CYCLES)

diff --git a/tests/integration/src/control_panel_tests.rs b/tests/integration/src/control_panel_tests.rs
@@ -445,10 +445,10 @@ fn deploy_too_many_stations() {
         assert!(res.0.is_ok());
 
         // deploying a new station should fail nonetheless
-        assert_eq!(
-            can_deploy(user_id).unwrap_err().code,
-            "DEPLOY_STATION_QUOTA_EXCEEDED"
-        );
+        assert!(matches!(
+            can_deploy(user_id).unwrap(),
+            CanDeployStationResponse::QuotaExceeded
+        ));
         assert_eq!(
             deploy(user_id, day, max_stations_per_user)
                 .unwrap_err()
@@ -467,12 +467,10 @@ fn deploy_too_many_stations() {
         }
 
         // deploying one more station on behalf of yet another use should fail due to global rate limit
-        assert_eq!(
-            can_deploy(user_test_id(max_stations_per_day))
-                .unwrap_err()
-                .code,
-            "DEPLOY_STATION_QUOTA_EXCEEDED"
-        );
+        assert!(matches!(
+            can_deploy(user_test_id(max_stations_per_day)).unwrap(),
+            CanDeployStationResponse::QuotaExceeded
+        ));
         assert_eq!(
             deploy(user_test_id(max_stations_per_day), day, 0)
                 .unwrap_err()