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