canic_testkit/
pic.rs

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