test and canister interface change

dfinity · Oct 28, 2024 · 279595e · 279595e
1 parent 4e14a6f
commit 279595e
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 6 deletions.
diff --git a/e2e/tests-dfx/assetscanister.bash b/e2e/tests-dfx/assetscanister.bash
@@ -1988,6 +1988,45 @@ WARN: {
   assert_command      dfx deploy
 }
 
+@test "set permissions through init argument" {
+  dfx_start
+  dfx deploy
+
+  assert_command dfx canister call e2e_project_frontend list_permitted '(record { permission = variant { Prepare }; })'
+  assert_eq "(vec {})"
+  assert_command dfx canister call e2e_project_frontend list_permitted '(record { permission = variant { Commit }; })'
+  assert_match "$(dfx identity get-principal)"
+  assert_command dfx canister call e2e_project_frontend list_permitted '(record { permission = variant { ManagePermissions }; })'
+  assert_eq "(vec {})"
+
+  dfx identity new alice --storage-mode plaintext
+  ALICE="$(dfx --identity alice identity get-principal)"
+
+  dfx canister install e2e_project_frontend --mode reinstall --argument "(opt variant {
+    Upgrade = record {
+      set_permissions = opt record {
+        prepare = vec {
+          principal \"${ALICE}\";
+        };
+        commit = vec {
+          principal \"$(dfx identity get-principal)\";
+          principal \"aaaaa-aa\";
+        };
+        manage_permissions = vec {
+          principal \"$(dfx identity get-principal)\";
+        };
+      }
+    }
+  })"
+  assert_command dfx canister call e2e_project_frontend list_permitted '(record { permission = variant { Prepare }; })'
+  assert_match "${ALICE}"
+  assert_command dfx canister call e2e_project_frontend list_permitted '(record { permission = variant { Commit }; })'
+  assert_match "$(dfx identity get-principal)"
+  assert_match '"aaaaa-aa"'
+  assert_command dfx canister call e2e_project_frontend list_permitted '(record { permission = variant { ManagePermissions }; })'
+  assert_match "$(dfx identity get-principal)"
+}
+
 @test "set permissions through upgrade argument" {
   dfx_start
   dfx deploy

diff --git a/src/canisters/frontend/ic-certified-assets/src/lib.rs b/src/canisters/frontend/ic-certified-assets/src/lib.rs
@@ -423,16 +423,23 @@ fn is_controller() -> Result<(), String> {
 }
 
 pub fn init(args: Option<AssetCanisterArgs>) {
-    if let Some(upgrade_arg) = args {
-        let AssetCanisterArgs::Init(InitArgs {}) = upgrade_arg else {
-            ic_cdk::trap("Cannot initialize the canister with an Upgrade argument. Please provide an Init argument.")
-        };
-    }
     STATE.with(|s| {
         let mut s = s.borrow_mut();
         s.clear();
         s.grant_permission(caller(), &Permission::Commit);
     });
+
+    if let Some(upgrade_arg) = args {
+        let AssetCanisterArgs::Init(init_args) = upgrade_arg else {
+            ic_cdk::trap("Cannot initialize the canister with an Upgrade argument. Please provide an Init argument.")
+        };
+        STATE.with(|s| {
+            let mut state = s.borrow_mut();
+            if let Some(set_permissions) = init_args.set_permissions {
+                state.set_permissions(set_permissions);
+            }
+        });
+    }
 }
 
 pub fn pre_upgrade() -> StableState {

diff --git a/src/canisters/frontend/ic-certified-assets/src/types.rs b/src/canisters/frontend/ic-certified-assets/src/types.rs
@@ -207,7 +207,9 @@ pub enum AssetCanisterArgs {
 }
 
 #[derive(Clone, Debug, CandidType, Deserialize)]
-pub struct InitArgs {}
+pub struct InitArgs {
+    pub set_permissions: Option<SetPermissions>,
+}
 
 #[derive(Clone, Debug, CandidType, Deserialize)]
 pub struct UpgradeArgs {