Skip to main content

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