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 Pic {
104    /// Install a root canister with the default root init arguments.
105    pub fn create_and_install_root_canister(&self, wasm: Vec<u8>) -> Result<Principal, Error> {
106        let init_bytes = install_root_args()?;
107
108        Ok(self.create_funded_and_install(wasm, init_bytes))
109    }
110
111    /// Install a canister with the given type and wasm bytes.
112    ///
113    /// Install failures are treated as fatal in tests.
114    pub fn create_and_install_canister(
115        &self,
116        role: CanisterRole,
117        wasm: Vec<u8>,
118    ) -> Result<Principal, Error> {
119        let init_bytes = install_args(role)?;
120
121        Ok(self.create_funded_and_install(wasm, init_bytes))
122    }
123
124    fn create_funded_and_install(&self, wasm: Vec<u8>, init_bytes: Vec<u8>) -> Principal {
125        let canister_id = self.create_canister();
126        self.add_cycles(canister_id, INSTALL_CYCLES);
127        self.inner
128            .install_canister(canister_id, wasm, init_bytes, None);
129
130        canister_id
131    }
132
133    /// Generic update call helper (serializes args + decodes result).
134    pub fn update_call<T, A>(
135        &self,
136        canister_id: Principal,
137        method: &str,
138        args: A,
139    ) -> Result<T, Error>
140    where
141        T: CandidType + DeserializeOwned,
142        A: ArgumentEncoder,
143    {
144        let bytes: Vec<u8> = encode_args(args)
145            .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
146        let result = self
147            .inner
148            .update_call(canister_id, Principal::anonymous(), method, bytes)
149            .map_err(|err| {
150                Error::internal(format!(
151                    "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
152                ))
153            })?;
154
155        decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
156    }
157
158    /// Generic update call helper with an explicit caller principal.
159    pub fn update_call_as<T, A>(
160        &self,
161        canister_id: Principal,
162        caller: Principal,
163        method: &str,
164        args: A,
165    ) -> Result<T, Error>
166    where
167        T: CandidType + DeserializeOwned,
168        A: ArgumentEncoder,
169    {
170        let bytes: Vec<u8> = encode_args(args)
171            .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
172        let result = self
173            .inner
174            .update_call(canister_id, caller, method, bytes)
175            .map_err(|err| {
176                Error::internal(format!(
177                    "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
178                ))
179            })?;
180
181        decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
182    }
183
184    /// Generic query call helper.
185    pub fn query_call<T, A>(
186        &self,
187        canister_id: Principal,
188        method: &str,
189        args: A,
190    ) -> Result<T, Error>
191    where
192        T: CandidType + DeserializeOwned,
193        A: ArgumentEncoder,
194    {
195        let bytes: Vec<u8> = encode_args(args)
196            .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
197        let result = self
198            .inner
199            .query_call(canister_id, Principal::anonymous(), method, bytes)
200            .map_err(|err| {
201                Error::internal(format!(
202                    "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
203                ))
204            })?;
205
206        decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
207    }
208
209    /// Generic query call helper with an explicit caller principal.
210    pub fn query_call_as<T, A>(
211        &self,
212        canister_id: Principal,
213        caller: Principal,
214        method: &str,
215        args: A,
216    ) -> Result<T, Error>
217    where
218        T: CandidType + DeserializeOwned,
219        A: ArgumentEncoder,
220    {
221        let bytes: Vec<u8> = encode_args(args)
222            .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
223        let result = self
224            .inner
225            .query_call(canister_id, caller, method, bytes)
226            .map_err(|err| {
227                Error::internal(format!(
228                    "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
229                ))
230            })?;
231
232        decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
233    }
234
235    pub fn tick_n(&self, times: usize) {
236        for _ in 0..times {
237            self.tick();
238        }
239    }
240}
241
242impl Deref for Pic {
243    type Target = PocketIc;
244
245    fn deref(&self) -> &Self::Target {
246        &self.inner
247    }
248}
249
250impl DerefMut for Pic {
251    fn deref_mut(&mut self) -> &mut Self::Target {
252        &mut self.inner
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}