canic_testkit/
pic.rs

1use candid::{CandidType, decode_one, encode_args, encode_one, utils::ArgumentEncoder};
2use canic::{
3    Error,
4    core::{ids::CanisterRole, model::memory::topology::SubnetIdentity, ops::CanisterInitPayload},
5    types::Principal,
6};
7use derive_more::{Deref, DerefMut};
8use pocket_ic::{PocketIc, PocketIcBuilder};
9use serde::de::DeserializeOwned;
10
11///
12/// PicBuilder
13///
14
15pub struct PicBuilder(PocketIcBuilder);
16
17#[allow(clippy::new_without_default)]
18impl PicBuilder {
19    /// Start a new PicBuilder with sensible defaults
20    #[must_use]
21    pub fn new() -> Self {
22        Self(PocketIcBuilder::new())
23    }
24
25    #[must_use]
26    pub fn with_application_subnet(mut self) -> Self {
27        self.0 = self.0.with_application_subnet();
28        self
29    }
30
31    #[must_use]
32    pub fn with_nns_subnet(mut self) -> Self {
33        self.0 = self.0.with_nns_subnet();
34        self
35    }
36
37    /// Finish building the PocketIC instance and wrap it
38    #[must_use]
39    pub fn build(self) -> Pic {
40        Pic(self.0.build())
41    }
42}
43
44///
45/// Pic
46///
47
48#[derive(Deref, DerefMut)]
49pub struct Pic(PocketIc);
50
51impl Pic {
52    /// Install a canister with the given type and wasm bytes
53    pub fn create_and_install_canister(
54        &self,
55        ty: CanisterRole,
56        wasm: Vec<u8>,
57    ) -> Result<Principal, Error> {
58        // Create and fund the canister
59        let canister_id = self.create_canister();
60        self.add_cycles(canister_id, 1_000_000_000_000);
61
62        // Install
63        let init_bytes = install_args(ty)?;
64        self.0.install_canister(canister_id, wasm, init_bytes, None);
65
66        Ok(canister_id)
67    }
68
69    /// Generic update call helper (serializes args + decodes result)
70    pub fn update_call<T, A>(
71        &self,
72        canister_id: Principal,
73        method: &str,
74        args: A,
75    ) -> Result<T, Error>
76    where
77        T: CandidType + DeserializeOwned,
78        A: ArgumentEncoder,
79    {
80        let bytes: Vec<u8> = encode_args(args)?;
81        let result = self
82            .0
83            .update_call(canister_id, Principal::anonymous(), method, bytes)
84            .map_err(|e| Error::test(e.to_string()))?;
85
86        decode_one(&result).map_err(Into::into)
87    }
88
89    /// Generic query call helper
90    pub fn query_call<T, A>(
91        &self,
92        canister_id: Principal,
93        method: &str,
94        args: A,
95    ) -> Result<T, Error>
96    where
97        T: CandidType + DeserializeOwned,
98        A: ArgumentEncoder,
99    {
100        let bytes: Vec<u8> = encode_args(args)?;
101        let result = self
102            .0
103            .query_call(canister_id, Principal::anonymous(), method, bytes)
104            .map_err(|e| Error::test(e.to_string()))?;
105
106        decode_one(&result).map_err(Into::into)
107    }
108}
109
110/// --------------------------------------
111/// install_args helper
112/// --------------------------------------
113fn install_args(ty: CanisterRole) -> Result<Vec<u8>, Error> {
114    let args = if ty.is_root() {
115        // Provide a deterministic subnet principal for PocketIC runs
116        let subnet_pid = Principal::from_slice(&[0xAA; 29]);
117        encode_one(SubnetIdentity::Manual(subnet_pid))
118    } else {
119        let payload = CanisterInitPayload::empty();
120        encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
121    }?;
122
123    Ok(args)
124}