canic_testkit/
pic.rs

1use candid::{CandidType, Principal, decode_one, encode_args, encode_one, utils::ArgumentEncoder};
2use canic::{
3    PublicError,
4    cdk::types::TC,
5    core::{
6        dto::{
7            abi::v1::CanisterInitPayload,
8            env::EnvView,
9            subnet::SubnetIdentity,
10            topology::{AppDirectoryView, SubnetDirectoryView},
11        },
12        ids::CanisterRole,
13    },
14};
15use derive_more::{Deref, DerefMut};
16use pocket_ic::{PocketIc, PocketIcBuilder};
17use serde::de::DeserializeOwned;
18
19const INSTALL_CYCLES: u128 = 500 * TC;
20
21///
22/// Create a fresh PocketIC universe.
23///
24/// IMPORTANT:
25/// - Each call creates a new IC instance
26/// - This must NOT be cached or shared across tests
27/// - Required to avoid PocketIC wasm chunk store exhaustion
28///
29#[must_use]
30pub fn pic() -> Pic {
31    PicBuilder::new().with_application_subnet().build()
32}
33
34///
35/// PicBuilder
36/// Thin wrapper around the PocketIC builder.
37///
38/// This builder is only used to configure the singleton. It does not create
39/// additional IC instances beyond the global `Pic`.
40///
41/// Note: this file is test-only infrastructure; simplicity wins over abstraction.
42///
43pub struct PicBuilder(PocketIcBuilder);
44
45#[allow(clippy::new_without_default)]
46impl PicBuilder {
47    /// Start a new PicBuilder with sensible defaults.
48    #[must_use]
49    pub fn new() -> Self {
50        Self(PocketIcBuilder::new())
51    }
52
53    /// Include an application subnet in the singleton universe.
54    #[must_use]
55    pub fn with_application_subnet(mut self) -> Self {
56        self.0 = self.0.with_application_subnet();
57        self
58    }
59
60    /// Include an NNS subnet in the singleton universe.
61    #[must_use]
62    pub fn with_nns_subnet(mut self) -> Self {
63        self.0 = self.0.with_nns_subnet();
64        self
65    }
66
67    /// Finish building the singleton PocketIC instance and wrap it.
68    #[must_use]
69    pub fn build(self) -> Pic {
70        Pic(self.0.build())
71    }
72}
73
74///
75/// Pic
76/// Thin wrapper around the global PocketIC instance.
77///
78/// This type intentionally exposes only a minimal API surface; callers should
79/// use `pic()` to obtain the singleton and then perform installs/calls.
80///
81#[derive(Deref, DerefMut)]
82pub struct Pic(PocketIc);
83
84impl Pic {
85    /// Install a root canister with the default root init arguments.
86    pub fn create_and_install_root_canister(
87        &self,
88        wasm: Vec<u8>,
89    ) -> Result<Principal, PublicError> {
90        let init_bytes = install_root_args()?;
91
92        Ok(self.create_funded_and_install(wasm, init_bytes))
93    }
94
95    /// Install a canister with the given type and wasm bytes.
96    ///
97    /// Install failures are treated as fatal in tests.
98    pub fn create_and_install_canister(
99        &self,
100        role: CanisterRole,
101        wasm: Vec<u8>,
102    ) -> Result<Principal, PublicError> {
103        let init_bytes = install_args(role)?;
104
105        Ok(self.create_funded_and_install(wasm, init_bytes))
106    }
107
108    /// Install a canister with a custom directory snapshot (local-only helper).
109    ///
110    /// Use this when a test exercises directory-dependent auth/endpoints and
111    /// cannot rely on root to provide a snapshot.
112    pub fn create_and_install_canister_with_directories(
113        &self,
114        role: CanisterRole,
115        wasm: Vec<u8>,
116        app_directory: AppDirectoryView,
117        subnet_directory: SubnetDirectoryView,
118    ) -> Result<Principal, PublicError> {
119        let init_bytes = install_args_with_directories(role, app_directory, subnet_directory)?;
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.0.install_canister(canister_id, wasm, init_bytes, None);
128
129        canister_id
130    }
131
132    /// Generic update call helper (serializes args + decodes result).
133    pub fn update_call<T, A>(
134        &self,
135        canister_id: Principal,
136        method: &str,
137        args: A,
138    ) -> Result<T, PublicError>
139    where
140        T: CandidType + DeserializeOwned,
141        A: ArgumentEncoder,
142    {
143        let bytes: Vec<u8> = encode_args(args)
144            .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))?;
145        let result = self
146            .0
147            .update_call(canister_id, Principal::anonymous(), method, bytes)
148            .map_err(|err| {
149                PublicError::internal(format!(
150                    "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
151                ))
152            })?;
153
154        decode_one(&result)
155            .map_err(|err| PublicError::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, PublicError>
166    where
167        T: CandidType + DeserializeOwned,
168        A: ArgumentEncoder,
169    {
170        let bytes: Vec<u8> = encode_args(args)
171            .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))?;
172        let result = self
173            .0
174            .update_call(canister_id, caller, method, bytes)
175            .map_err(|err| {
176                PublicError::internal(format!(
177                    "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
178                ))
179            })?;
180
181        decode_one(&result)
182            .map_err(|err| PublicError::internal(format!("decode_one failed: {err}")))
183    }
184
185    /// Generic query call helper.
186    pub fn query_call<T, A>(
187        &self,
188        canister_id: Principal,
189        method: &str,
190        args: A,
191    ) -> Result<T, PublicError>
192    where
193        T: CandidType + DeserializeOwned,
194        A: ArgumentEncoder,
195    {
196        let bytes: Vec<u8> = encode_args(args)
197            .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))?;
198        let result = self
199            .0
200            .query_call(canister_id, Principal::anonymous(), method, bytes)
201            .map_err(|err| {
202                PublicError::internal(format!(
203                    "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
204                ))
205            })?;
206
207        decode_one(&result)
208            .map_err(|err| PublicError::internal(format!("decode_one failed: {err}")))
209    }
210
211    /// Generic query call helper with an explicit caller principal.
212    pub fn query_call_as<T, A>(
213        &self,
214        canister_id: Principal,
215        caller: Principal,
216        method: &str,
217        args: A,
218    ) -> Result<T, PublicError>
219    where
220        T: CandidType + DeserializeOwned,
221        A: ArgumentEncoder,
222    {
223        let bytes: Vec<u8> = encode_args(args)
224            .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))?;
225        let result = self
226            .0
227            .query_call(canister_id, caller, method, bytes)
228            .map_err(|err| {
229                PublicError::internal(format!(
230                    "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
231                ))
232            })?;
233
234        decode_one(&result)
235            .map_err(|err| PublicError::internal(format!("decode_one failed: {err}")))
236    }
237
238    pub fn tick_n(&self, times: usize) {
239        for _ in 0..times {
240            self.tick();
241        }
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 `EnvView` + 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>, PublicError> {
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 = EnvView {
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: AppDirectoryView(Vec::new()),
278            subnet_directory: SubnetDirectoryView(Vec::new()),
279        };
280
281        encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
282            .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))
283    }
284}
285
286fn install_root_args() -> Result<Vec<u8>, PublicError> {
287    encode_one(SubnetIdentity::Manual)
288        .map_err(|err| PublicError::internal(format!("encode_one failed: {err}")))
289}
290
291fn install_args_with_directories(
292    role: CanisterRole,
293    app_directory: AppDirectoryView,
294    subnet_directory: SubnetDirectoryView,
295) -> Result<Vec<u8>, PublicError> {
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| PublicError::internal(format!("encode_one failed: {err}")))
301    } else {
302        // Non-root canister: pass structural context, not invented identities.
303        let env = EnvView {
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| PublicError::internal(format!("encode_args failed: {err}")))
320    }
321}