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