canic_testkit/
pic.rs

1use candid::{CandidType, Principal, decode_one, encode_args, encode_one, utils::ArgumentEncoder};
2use canic::core::{
3    PublicError,
4    cdk::types::TC,
5    dto::{
6        abi::v1::CanisterInitPayload,
7        env::EnvView,
8        subnet::SubnetIdentity,
9        topology::{AppDirectoryView, SubnetDirectoryView},
10    },
11    ids::CanisterRole,
12};
13use derive_more::{Deref, DerefMut};
14use pocket_ic::{PocketIc, PocketIcBuilder};
15use serde::de::DeserializeOwned;
16
17const INSTALL_CYCLES: u128 = 500 * TC;
18
19///
20/// Create a fresh PocketIC universe.
21///
22/// IMPORTANT:
23/// - Each call creates a new IC instance
24/// - This must NOT be cached or shared across tests
25/// - Required to avoid PocketIC wasm chunk store exhaustion
26///
27#[must_use]
28pub fn pic() -> Pic {
29    PicBuilder::new().with_application_subnet().build()
30}
31
32///
33/// PicBuilder
34/// Thin wrapper around the PocketIC builder.
35///
36/// This builder is only used to configure the singleton. It does not create
37/// additional IC instances beyond the global `Pic`.
38///
39/// Note: this file is test-only infrastructure; simplicity wins over abstraction.
40///
41pub struct PicBuilder(PocketIcBuilder);
42
43#[allow(clippy::new_without_default)]
44impl PicBuilder {
45    /// Start a new PicBuilder with sensible defaults.
46    #[must_use]
47    pub fn new() -> Self {
48        Self(PocketIcBuilder::new())
49    }
50
51    /// Include an application subnet in the singleton universe.
52    #[must_use]
53    pub fn with_application_subnet(mut self) -> Self {
54        self.0 = self.0.with_application_subnet();
55        self
56    }
57
58    /// Include an NNS subnet in the singleton universe.
59    #[must_use]
60    pub fn with_nns_subnet(mut self) -> Self {
61        self.0 = self.0.with_nns_subnet();
62        self
63    }
64
65    /// Finish building the singleton PocketIC instance and wrap it.
66    #[must_use]
67    pub fn build(self) -> Pic {
68        Pic(self.0.build())
69    }
70}
71
72///
73/// Pic
74/// Thin wrapper around the global PocketIC instance.
75///
76/// This type intentionally exposes only a minimal API surface; callers should
77/// use `pic()` to obtain the singleton and then perform installs/calls.
78///
79#[derive(Deref, DerefMut)]
80pub struct Pic(PocketIc);
81
82impl Pic {
83    /// Install a root canister with the default root init arguments.
84    pub fn create_and_install_root_canister(
85        &self,
86        wasm: Vec<u8>,
87    ) -> Result<Principal, PublicError> {
88        let init_bytes = install_root_args()?;
89
90        Ok(self.create_funded_and_install(wasm, init_bytes))
91    }
92
93    /// Install a canister with the given type and wasm bytes.
94    ///
95    /// Install failures are treated as fatal in tests.
96    pub fn create_and_install_canister(
97        &self,
98        role: CanisterRole,
99        wasm: Vec<u8>,
100    ) -> Result<Principal, PublicError> {
101        let init_bytes = install_args(role)?;
102
103        Ok(self.create_funded_and_install(wasm, init_bytes))
104    }
105
106    /// Install a canister with a custom directory snapshot (local-only helper).
107    ///
108    /// Use this when a test exercises directory-dependent auth/endpoints and
109    /// cannot rely on root to provide a snapshot.
110    pub fn create_and_install_canister_with_directories(
111        &self,
112        role: CanisterRole,
113        wasm: Vec<u8>,
114        app_directory: AppDirectoryView,
115        subnet_directory: SubnetDirectoryView,
116    ) -> Result<Principal, PublicError> {
117        let init_bytes = install_args_with_directories(role, app_directory, subnet_directory)?;
118
119        Ok(self.create_funded_and_install(wasm, init_bytes))
120    }
121
122    fn create_funded_and_install(&self, wasm: Vec<u8>, init_bytes: Vec<u8>) -> Principal {
123        let canister_id = self.create_canister();
124        self.add_cycles(canister_id, INSTALL_CYCLES);
125        self.0.install_canister(canister_id, wasm, init_bytes, None);
126
127        canister_id
128    }
129
130    /// Generic update call helper (serializes args + decodes result).
131    pub fn update_call<T, A>(
132        &self,
133        canister_id: Principal,
134        method: &str,
135        args: A,
136    ) -> Result<T, PublicError>
137    where
138        T: CandidType + DeserializeOwned,
139        A: ArgumentEncoder,
140    {
141        let bytes: Vec<u8> = encode_args(args)
142            .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))?;
143        let result = self
144            .0
145            .update_call(canister_id, Principal::anonymous(), method, bytes)
146            .map_err(|err| {
147                PublicError::internal(format!(
148                    "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
149                ))
150            })?;
151
152        decode_one(&result)
153            .map_err(|err| PublicError::internal(format!("decode_one failed: {err}")))
154    }
155
156    /// Generic update call helper with an explicit caller principal.
157    pub fn update_call_as<T, A>(
158        &self,
159        canister_id: Principal,
160        caller: Principal,
161        method: &str,
162        args: A,
163    ) -> Result<T, PublicError>
164    where
165        T: CandidType + DeserializeOwned,
166        A: ArgumentEncoder,
167    {
168        let bytes: Vec<u8> = encode_args(args)
169            .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))?;
170        let result = self
171            .0
172            .update_call(canister_id, caller, method, bytes)
173            .map_err(|err| {
174                PublicError::internal(format!(
175                    "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
176                ))
177            })?;
178
179        decode_one(&result)
180            .map_err(|err| PublicError::internal(format!("decode_one failed: {err}")))
181    }
182
183    /// Generic query call helper.
184    pub fn query_call<T, A>(
185        &self,
186        canister_id: Principal,
187        method: &str,
188        args: A,
189    ) -> Result<T, PublicError>
190    where
191        T: CandidType + DeserializeOwned,
192        A: ArgumentEncoder,
193    {
194        let bytes: Vec<u8> = encode_args(args)
195            .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))?;
196        let result = self
197            .0
198            .query_call(canister_id, Principal::anonymous(), method, bytes)
199            .map_err(|err| {
200                PublicError::internal(format!(
201                    "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
202                ))
203            })?;
204
205        decode_one(&result)
206            .map_err(|err| PublicError::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, PublicError>
217    where
218        T: CandidType + DeserializeOwned,
219        A: ArgumentEncoder,
220    {
221        let bytes: Vec<u8> = encode_args(args)
222            .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))?;
223        let result = self
224            .0
225            .query_call(canister_id, caller, method, bytes)
226            .map_err(|err| {
227                PublicError::internal(format!(
228                    "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
229                ))
230            })?;
231
232        decode_one(&result)
233            .map_err(|err| PublicError::internal(format!("decode_one failed: {err}")))
234    }
235
236    pub fn tick_n(&self, times: usize) {
237        for _ in 0..times {
238            self.tick();
239        }
240    }
241}
242
243/// --------------------------------------
244/// install_args helper
245/// --------------------------------------
246///
247/// Init semantics:
248/// - Root canisters receive a `SubnetIdentity` (direct root bootstrap).
249/// - Non-root canisters receive `EnvView` + optional directory snapshots.
250///
251/// Directory handling:
252/// - By default, directory views are empty for standalone installs.
253/// - Directory-dependent logic is opt-in via `install_args_with_directories`.
254/// - Root-provisioned installs will populate directories via cascade.
255///
256fn install_args(role: CanisterRole) -> Result<Vec<u8>, PublicError> {
257    if role.is_root() {
258        install_root_args()
259    } else {
260        // Non-root standalone install.
261        // Provide only what is structurally known at install time.
262        let env = EnvView {
263            prime_root_pid: None,
264            subnet_role: None,
265            subnet_pid: None,
266            root_pid: None,
267            canister_role: Some(role),
268            parent_pid: None,
269        };
270
271        // Intentional: standalone installs do not require directories unless
272        // a test explicitly exercises directory-dependent behavior.
273        let payload = CanisterInitPayload {
274            env,
275            app_directory: AppDirectoryView(Vec::new()),
276            subnet_directory: SubnetDirectoryView(Vec::new()),
277        };
278
279        encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
280            .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))
281    }
282}
283
284fn install_root_args() -> Result<Vec<u8>, PublicError> {
285    encode_one(SubnetIdentity::Manual)
286        .map_err(|err| PublicError::internal(format!("encode_one failed: {err}")))
287}
288
289fn install_args_with_directories(
290    role: CanisterRole,
291    app_directory: AppDirectoryView,
292    subnet_directory: SubnetDirectoryView,
293) -> Result<Vec<u8>, PublicError> {
294    if role.is_root() {
295        // Root canister: runtime identity only.
296        // No fake principals. Runtime/bootstrap will resolve actual context.
297        encode_one(SubnetIdentity::Manual)
298            .map_err(|err| PublicError::internal(format!("encode_one failed: {err}")))
299    } else {
300        // Non-root canister: pass structural context, not invented identities.
301        let env = EnvView {
302            prime_root_pid: None,
303            subnet_role: None,
304            subnet_pid: None,
305            root_pid: None,
306            canister_role: Some(role),
307            parent_pid: None,
308        };
309
310        let payload = CanisterInitPayload {
311            env,
312            app_directory,
313            subnet_directory,
314        };
315
316        encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
317            .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))
318    }
319}
canic_testkit/pic.rs

canic_testkit/
pic.rs
