canic_testkit/
pic.rs

1use candid::{CandidType, Principal, decode_one, encode_args, encode_one, utils::ArgumentEncoder};
2use canic::{
3    Error,
4    cdk::types::TC,
5    dto::{
6        abi::v1::CanisterInitPayload,
7        env::EnvBootstrapArgs,
8        subnet::SubnetIdentity,
9        topology::{AppDirectoryArgs, SubnetDirectoryArgs},
10    },
11    ids::CanisterRole,
12};
13use derive_more::{Deref, DerefMut};
14use pocket_ic::{PocketIc, PocketIcBuilder};
15use serde::de::DeserializeOwned;
16use std::time::Duration;
17
18const INSTALL_CYCLES: u128 = 500 * TC;
19
20///
21/// Create a fresh PocketIC universe.
22///
23/// IMPORTANT:
24/// - Each call creates a new IC instance
25/// - This must NOT be cached or shared across tests
26/// - Required to avoid PocketIC wasm chunk store exhaustion
27///
28#[must_use]
29pub fn pic() -> Pic {
30    PicBuilder::new().with_application_subnet().build()
31}
32
33///
34/// PicBuilder
35/// Thin wrapper around the PocketIC builder.
36///
37/// This builder is only used to configure the singleton. It does not create
38/// additional IC instances beyond the global `Pic`.
39///
40/// Note: this file is test-only infrastructure; simplicity wins over abstraction.
41///
42pub struct PicBuilder(PocketIcBuilder);
43
44#[allow(clippy::new_without_default)]
45impl PicBuilder {
46    /// Start a new PicBuilder with sensible defaults.
47    #[must_use]
48    pub fn new() -> Self {
49        Self(PocketIcBuilder::new())
50    }
51
52    /// Include an application subnet in the singleton universe.
53    #[must_use]
54    pub fn with_application_subnet(mut self) -> Self {
55        self.0 = self.0.with_application_subnet();
56        self
57    }
58
59    /// Include an NNS subnet in the singleton universe.
60    #[must_use]
61    pub fn with_nns_subnet(mut self) -> Self {
62        self.0 = self.0.with_nns_subnet();
63        self
64    }
65
66    /// Finish building the singleton PocketIC instance and wrap it.
67    #[must_use]
68    pub fn build(self) -> Pic {
69        Pic(self.0.build())
70    }
71}
72
73///
74/// Pic
75/// Thin wrapper around the global PocketIC instance.
76///
77/// This type intentionally exposes only a minimal API surface; callers should
78/// use `pic()` to obtain the singleton and then perform installs/calls.
79///
80#[derive(Deref, DerefMut)]
81pub struct Pic(PocketIc);
82
83impl Pic {
84    /// Install a root canister with the default root init arguments.
85    pub fn create_and_install_root_canister(&self, wasm: Vec<u8>) -> Result<Principal, Error> {
86        let init_bytes = install_root_args()?;
87
88        Ok(self.create_funded_and_install(wasm, init_bytes))
89    }
90
91    /// Install a canister with the given type and wasm bytes.
92    ///
93    /// Install failures are treated as fatal in tests.
94    pub fn create_and_install_canister(
95        &self,
96        role: CanisterRole,
97        wasm: Vec<u8>,
98    ) -> Result<Principal, Error> {
99        let init_bytes = install_args(role)?;
100
101        Ok(self.create_funded_and_install(wasm, init_bytes))
102    }
103
104    /// Install a canister with a custom directory snapshot (local-only helper).
105    ///
106    /// Use this when a test exercises directory-dependent auth/endpoints and
107    /// cannot rely on root to provide a snapshot.
108    pub fn create_and_install_canister_with_directories(
109        &self,
110        role: CanisterRole,
111        wasm: Vec<u8>,
112        app_directory: AppDirectoryArgs,
113        subnet_directory: SubnetDirectoryArgs,
114    ) -> Result<Principal, Error> {
115        let init_bytes = install_args_with_directories(role, app_directory, subnet_directory)?;
116
117        Ok(self.create_funded_and_install(wasm, init_bytes))
118    }
119
120    fn create_funded_and_install(&self, wasm: Vec<u8>, init_bytes: Vec<u8>) -> Principal {
121        let canister_id = self.create_canister();
122        self.add_cycles(canister_id, INSTALL_CYCLES);
123        self.0.install_canister(canister_id, wasm, init_bytes, None);
124
125        canister_id
126    }
127
128    /// Generic update call helper (serializes args + decodes result).
129    pub fn update_call<T, A>(
130        &self,
131        canister_id: Principal,
132        method: &str,
133        args: A,
134    ) -> Result<T, Error>
135    where
136        T: CandidType + DeserializeOwned,
137        A: ArgumentEncoder,
138    {
139        let bytes: Vec<u8> = encode_args(args)
140            .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
141        let result = self
142            .0
143            .update_call(canister_id, Principal::anonymous(), method, bytes)
144            .map_err(|err| {
145                Error::internal(format!(
146                    "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
147                ))
148            })?;
149
150        decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
151    }
152
153    /// Generic update call helper with an explicit caller principal.
154    pub fn update_call_as<T, A>(
155        &self,
156        canister_id: Principal,
157        caller: Principal,
158        method: &str,
159        args: A,
160    ) -> Result<T, Error>
161    where
162        T: CandidType + DeserializeOwned,
163        A: ArgumentEncoder,
164    {
165        let bytes: Vec<u8> = encode_args(args)
166            .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
167        let result = self
168            .0
169            .update_call(canister_id, caller, method, bytes)
170            .map_err(|err| {
171                Error::internal(format!(
172                    "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
173                ))
174            })?;
175
176        decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
177    }
178
179    /// Generic query call helper.
180    pub fn query_call<T, A>(
181        &self,
182        canister_id: Principal,
183        method: &str,
184        args: A,
185    ) -> Result<T, Error>
186    where
187        T: CandidType + DeserializeOwned,
188        A: ArgumentEncoder,
189    {
190        let bytes: Vec<u8> = encode_args(args)
191            .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
192        let result = self
193            .0
194            .query_call(canister_id, Principal::anonymous(), method, bytes)
195            .map_err(|err| {
196                Error::internal(format!(
197                    "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
198                ))
199            })?;
200
201        decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
202    }
203
204    /// Generic query call helper with an explicit caller principal.
205    pub fn query_call_as<T, A>(
206        &self,
207        canister_id: Principal,
208        caller: Principal,
209        method: &str,
210        args: A,
211    ) -> Result<T, Error>
212    where
213        T: CandidType + DeserializeOwned,
214        A: ArgumentEncoder,
215    {
216        let bytes: Vec<u8> = encode_args(args)
217            .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
218        let result = self
219            .0
220            .query_call(canister_id, caller, method, bytes)
221            .map_err(|err| {
222                Error::internal(format!(
223                    "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
224                ))
225            })?;
226
227        decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
228    }
229
230    pub fn tick_n(&self, times: usize) {
231        for _ in 0..times {
232            self.tick();
233        }
234    }
235
236    pub fn certify_time(&self) {
237        let now = self.0.get_time();
238        let next = now + Duration::from_secs(1);
239        self.0.set_time(next);
240        self.0.set_certified_time(next);
241        self.0.tick();
242    }
243}
244
245/// --------------------------------------
246/// install_args helper
247/// --------------------------------------
248///
249/// Init semantics:
250/// - Root canisters receive a `SubnetIdentity` (direct root bootstrap).
251/// - Non-root canisters receive `EnvBootstrapArgs` + optional directory snapshots.
252///
253/// Directory handling:
254/// - By default, directory views are empty for standalone installs.
255/// - Directory-dependent logic is opt-in via `install_args_with_directories`.
256/// - Root-provisioned installs will populate directories via cascade.
257///
258fn install_args(role: CanisterRole) -> Result<Vec<u8>, Error> {
259    if role.is_root() {
260        install_root_args()
261    } else {
262        // Non-root standalone install.
263        // Provide only what is structurally known at install time.
264        let env = EnvBootstrapArgs {
265            prime_root_pid: None,
266            subnet_role: None,
267            subnet_pid: None,
268            root_pid: None,
269            canister_role: Some(role),
270            parent_pid: None,
271        };
272
273        // Intentional: standalone installs do not require directories unless
274        // a test explicitly exercises directory-dependent behavior.
275        let payload = CanisterInitPayload {
276            env,
277            app_directory: AppDirectoryArgs(Vec::new()),
278            subnet_directory: SubnetDirectoryArgs(Vec::new()),
279        };
280
281        encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
282            .map_err(|err| Error::internal(format!("encode_args failed: {err}")))
283    }
284}
285
286fn install_root_args() -> Result<Vec<u8>, Error> {
287    encode_one(SubnetIdentity::Manual)
288        .map_err(|err| Error::internal(format!("encode_one failed: {err}")))
289}
290
291fn install_args_with_directories(
292    role: CanisterRole,
293    app_directory: AppDirectoryArgs,
294    subnet_directory: SubnetDirectoryArgs,
295) -> Result<Vec<u8>, Error> {
296    if role.is_root() {
297        // Root canister: runtime identity only.
298        // No fake principals. Runtime/bootstrap will resolve actual context.
299        encode_one(SubnetIdentity::Manual)
300            .map_err(|err| Error::internal(format!("encode_one failed: {err}")))
301    } else {
302        // Non-root canister: pass structural context, not invented identities.
303        let env = EnvBootstrapArgs {
304            prime_root_pid: None,
305            subnet_role: None,
306            subnet_pid: None,
307            root_pid: None,
308            canister_role: Some(role),
309            parent_pid: None,
310        };
311
312        let payload = CanisterInitPayload {
313            env,
314            app_directory,
315            subnet_directory,
316        };
317
318        encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
319            .map_err(|err| Error::internal(format!("encode_args failed: {err}")))
320    }
321}
canic_testkit/pic.rs

canic_testkit/
pic.rs
