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