refactor(nervous-system-agent): Add pocket-ic support to ic-nervous-system-agent

rs/nervous_system/agent/BUILD.bazel

@@ -4,6 +4,7 @@ package(default_visibility = ["//visibility:public"])
 
 DEPENDENCIES = [
     # Keep sorted.
+    "//packages/pocket-ic",
     "//rs/nervous_system/clients",
     "//rs/nns/constants",
     "//rs/nns/sns-wasm",

rs/nervous_system/agent/Cargo.toml

@@ -16,6 +16,7 @@ ic-nervous-system-clients = { path = "../clients" }
 ic-nns-constants = { path = "../../nns/constants" }
 ic-sns-wasm = { path = "../../nns/sns-wasm" }
 ic-sns-governance = { path = "../../sns/governance" }
+pocket-ic = { path = "../../../packages/pocket-ic" }
 serde = { workspace = true }
 tempfile = { workspace = true }
 thiserror = { workspace = true }

rs/nervous_system/agent/src/agent_impl.rs

@@ -0,0 +1,50 @@
+use candid::Principal;
+use ic_agent::Agent;
+use ic_nervous_system_clients::Request;
+use thiserror::Error;
+
+use crate::CallCanisters;
+
+#[derive(Error, Debug)]
+pub enum AgentCallError {
+    #[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),
+}
+
+impl CallCanisters for Agent {
+    type Error = AgentCallError;
+    async fn call<R: Request>(
+        &self,
+        canister_id: impl Into<Principal> + Send,
+        request: R,
+    ) -> Result<R::Response, Self::Error> {
+        let canister_id = canister_id.into();
+        let request_bytes = candid::encode_one(&request).map_err(AgentCallError::CandidEncode)?;
+        let response = if R::UPDATE {
+            let request = self
+                .update(&canister_id, R::METHOD)
+                .with_arg(request_bytes)
+                .call()
+                .await?;
+            match request {
+                ic_agent::agent::CallResponse::Response(response) => response,
+                ic_agent::agent::CallResponse::Poll(request_id) => {
+                    self.wait(&request_id, canister_id).await?
+                }
+            }
+        } else {
+            self.query(&canister_id, R::METHOD)
+                .with_arg(request_bytes)
+                .call()
+                .await?
+        };
+
+        let response =
+            candid::decode_one(response.as_slice()).map_err(AgentCallError::CandidDecode)?;
+        Ok(response)
+    }
+}

rs/nervous_system/agent/src/lib.rs

@@ -1,11 +1,11 @@
+mod agent_impl;
 pub mod nns;
+mod pocketic_impl;
 pub mod sns;
 
 use candid::Principal;
-use ic_agent::Agent;
 use ic_nervous_system_clients::Request;
 use std::fmt::Display;
-use thiserror::Error;
 
 pub trait CallCanisters {
     type Error: Display + Send;
@@ -15,48 +15,3 @@ pub trait CallCanisters {
         request: R,
     ) -> impl std::future::Future<Output = Result<R::Response, Self::Error>> + Send;
 }
-
-#[derive(Error, Debug)]
-pub enum AgentCallError {
-    #[error("agent error: {0}")]
-    AgentError(#[from] ic_agent::AgentError),
-    #[error("canister request could not be encoded: {0}")]
-    CandidEncodeError(candid::Error),
-    #[error("canister did not respond with the expected response type: {0}")]
-    CandidDecodeError(candid::Error),
-}
-
-impl CallCanisters for Agent {
-    type Error = AgentCallError;
-    async fn call<R: Request>(
-        &self,
-        canister_id: impl Into<Principal> + Send,
-        request: R,
-    ) -> Result<R::Response, Self::Error> {
-        let canister_id = canister_id.into();
-        let request_bytes =
-            candid::encode_one(&request).map_err(AgentCallError::CandidEncodeError)?;
-        let response = if R::UPDATE {
-            let request = self
-                .update(&canister_id, R::METHOD)
-                .with_arg(request_bytes)
-                .call()
-                .await?;
-            match request {
-                ic_agent::agent::CallResponse::Response(response) => response,
-                ic_agent::agent::CallResponse::Poll(request_id) => {
-                    self.wait(&request_id, canister_id).await?
-                }
-            }
-        } else {
-            self.query(&canister_id, R::METHOD)
-                .with_arg(request_bytes)
-                .call()
-                .await?
-        };
-
-        let response =
-            candid::decode_one(response.as_slice()).map_err(AgentCallError::CandidDecodeError)?;
-        Ok(response)
-    }
-}

rs/nervous_system/agent/src/pocketic_impl.rs

@@ -0,0 +1,56 @@
+use candid::Principal;
+use ic_nervous_system_clients::Request;
+use pocket_ic::PocketIc;
+use thiserror::Error;
+
+use crate::CallCanisters;
+
+#[derive(Error, Debug)]
+pub enum PocketIcCallError {
+    #[error("pocket_ic error: {0}")]
+    PocketIc(pocket_ic::UserError),
+    #[error("canister rejected the request: {0}")]
+    Reject(String),
+    #[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),
+}
+
+impl CallCanisters for PocketIc {
+    type Error = PocketIcCallError;
+    async fn call<R: Request>(
+        &self,
+        canister_id: impl Into<Principal> + Send,
+        request: R,
+    ) -> Result<R::Response, Self::Error> {
+        let canister_id = canister_id.into();
+        let request_bytes =
+            candid::encode_one(&request).map_err(PocketIcCallError::CandidEncode)?;
+        let response = if R::UPDATE {
+            self.update_call(
+                canister_id,
+                Principal::anonymous(),
+                R::METHOD,
+                request_bytes,
+            )
+        } else {
+            self.query_call(
+                canister_id,
+                Principal::anonymous(),
+                R::METHOD,
+                request_bytes,
+            )
+        }
+        .map_err(PocketIcCallError::PocketIc)?;
+
+        match response {
+            pocket_ic::WasmResult::Reply(reply) => {
+                let response = candid::decode_one(reply.as_slice())
+                    .map_err(PocketIcCallError::CandidDecode)?;
+                Ok(response)
+            }
+            pocket_ic::WasmResult::Reject(reject) => Err(PocketIcCallError::Reject(reject)),
+        }
+    }
+}