test(sns-cli): Port UpgradeSnsControlledCanister with Large Wasm integration test to use SNS CLI with PocketIc

packages/pocket-ic/src/nonblocking.rs

@@ -1529,7 +1529,7 @@ impl PocketIc {
         result.into()
     }
 
-    pub(crate) async fn update_call_with_effective_principal(
+    pub async fn update_call_with_effective_principal(
         &self,
         canister_id: CanisterId,
         effective_principal: RawEffectivePrincipal,

rs/nervous_system/agent/BUILD.bazel

@@ -19,7 +19,9 @@ DEPENDENCIES = [
     "@crate_index//:anyhow",
     "@crate_index//:candid",
     "@crate_index//:ic-agent",
+    "@crate_index//:itertools",
     "@crate_index//:serde",
+    "@crate_index//:serde_cbor",
     "@crate_index//:tempfile",
     "@crate_index//:thiserror",
     "@crate_index//:tokio",

rs/nervous_system/agent/Cargo.toml

@@ -23,7 +23,9 @@ pocket-ic = { path = "../../../packages/pocket-ic" }
 registry-canister = { path = "../../registry/canister" }
 ic-sns-root = { path = "../../sns/root" }
 ic-sns-swap = { path = "../../sns/swap" }
+itertools = { workspace = true }
 serde = { workspace = true }
+serde_cbor = { workspace = true }
 tempfile = { workspace = true }
 thiserror = { workspace = true }
 tokio = { workspace = true }

rs/nervous_system/agent/src/agent_impl.rs

@@ -1,24 +1,31 @@
-use crate::Request;
+use crate::CallCanisters;
+use crate::{CanisterInfo, Request};
 use candid::Principal;
 use ic_agent::Agent;
+use itertools::{Either, Itertools};
+use serde_cbor::Value;
+use std::collections::BTreeSet;
 use thiserror::Error;
 
-use crate::CallCanisters;
-
 #[derive(Error, Debug)]
 pub enum AgentCallError {
+    #[error("agent identity error: {0}")]
+    Identity(String),
     #[error("agent error: {0}")]
     Agent(#[from] ic_agent::AgentError),
     #[error("canister request could not be encoded: {0}")]
     CandidEncode(candid::Error),
     #[error("canister did not respond with the expected response type: {0}")]
     CandidDecode(candid::Error),
+    #[error("invalid canister controllers: {0}")]
+    CanisterControllers(String),
 }
 
 impl crate::sealed::Sealed for Agent {}
 
 impl CallCanisters for Agent {
     type Error = AgentCallError;
+
     async fn call<R: Request>(
         &self,
         canister_id: impl Into<Principal> + Send,
@@ -50,4 +57,86 @@ impl CallCanisters for Agent {
             candid::decode_one(response.as_slice()).map_err(AgentCallError::CandidDecode)?;
         Ok(response)
     }
+
+    async fn canister_info(
+        &self,
+        canister_id: impl Into<Principal> + Send,
+    ) -> Result<CanisterInfo, Self::Error> {
+        let canister_id = canister_id.into();
+
+        let read_state_result = self
+            .read_state_canister_info(canister_id, "module_hash")
+            .await;
+
+        let module_hash = match read_state_result {
+            Ok(module_hash) => Some(module_hash),
+            Err(ic_agent::AgentError::LookupPathAbsent(_)) => {
+                // https://internetcomputer.org/docs/current/references/ic-interface-spec#state-tree-canister-information
+                None
+            }
+            Err(err) => {
+                return Err(Self::Error::Agent(err));
+            }
+        };
+
+        let controllers_blob = self
+            .read_state_canister_info(canister_id, "controllers")
+            .await
+            .map_err(AgentCallError::Agent)?;
+
+        let cbor: Value = serde_cbor::from_slice(&controllers_blob).map_err(|err| {
+            Self::Error::CanisterControllers(format!("Failed decoding CBOR data: {:?}", err))
+        })?;
+
+        let Value::Array(controllers) = cbor else {
+            return Err(Self::Error::CanisterControllers(format!(
+                "Expected controllers to be an array, but got {:?}",
+                cbor
+            )));
+        };
+
+        let (controllers, errors): (Vec<_>, Vec<_>) =
+            controllers.into_iter().partition_map(|value| {
+                let Value::Bytes(bytes) = value else {
+                    let err = format!(
+                        "Expected canister controller to be of type bytes, got {:?}",
+                        value
+                    );
+                    return Either::Right(err);
+                };
+                match Principal::try_from(&bytes) {
+                    Err(err) => {
+                        let err =
+                            format!("Cannot interpret canister controller principal: {}", err);
+                        Either::Right(err)
+                    }
+                    Ok(principal) => Either::Left(principal),
+                }
+            });
+
+        if !errors.is_empty() {
+            return Err(Self::Error::CanisterControllers(format!(
+                "\n  - {}",
+                errors.join("\n  - ")
+            )));
+        }
+
+        let unique_controllers = BTreeSet::from_iter(controllers.iter().copied());
+
+        if unique_controllers.len() != controllers.len() {
+            return Err(Self::Error::CanisterControllers(format!(
+                "Canister controllers have duplicates: {}",
+                controllers.into_iter().join(", ")
+            )));
+        }
+
+        Ok(CanisterInfo {
+            module_hash,
+            controllers: unique_controllers,
+        })
+    }
+
+    fn caller(&self) -> Result<Principal, Self::Error> {
+        self.get_principal().map_err(Self::Error::Identity)
+    }
 }

rs/nervous_system/agent/src/lib.rs

@@ -1,14 +1,15 @@
+use candid::{CandidType, Principal};
+use serde::de::DeserializeOwned;
+use std::collections::BTreeSet;
+use std::fmt::Display;
+
 pub mod agent_impl;
 pub mod management_canister;
 pub mod nns;
 mod null_request;
 pub mod pocketic_impl;
 pub mod sns;
 
-use candid::{CandidType, Principal};
-use serde::de::DeserializeOwned;
-use std::fmt::Display;
-
 // This is used to "seal" the CallCanisters trait so that it cannot be implemented outside of this crate.
 // This is useful because it means we can modify the trait in the future without worrying about
 // breaking backwards compatibility with implementations outside of this crate.
@@ -27,25 +28,24 @@ pub trait Request: Send {
     type Response: CandidType + DeserializeOwned;
 }
 
-// impl<R: ic_nervous_system_clients::Request> Request for R {
-//     fn method(&self) -> &'static str {
-//         Self::METHOD
-//     }
-//     fn update(&self) -> bool {
-//         Self::UPDATE
-//     }
-//     fn payload(&self) -> Result<Vec<u8>, candid::Error> {
-//         candid::encode_one(self)
-//     }
-
-//     type Response = <Self as ic_nervous_system_clients::Request>::Response;
-// }
+pub struct CanisterInfo {
+    pub module_hash: Option<Vec<u8>>,
+    pub controllers: BTreeSet<Principal>,
+}
 
 pub trait CallCanisters: sealed::Sealed {
     type Error: Display + Send + std::error::Error + 'static;
+
+    fn caller(&self) -> Result<Principal, Self::Error>;
+
     fn call<R: Request>(
         &self,
         canister_id: impl Into<Principal> + Send,
         request: R,
     ) -> impl std::future::Future<Output = Result<R::Response, Self::Error>> + Send;
+
+    fn canister_info(
+        &self,
+        canister_id: impl Into<Principal> + Send,
+    ) -> impl std::future::Future<Output = Result<CanisterInfo, Self::Error>> + Send;
 }