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, SubnetRole},
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    let args = if role.is_root() {
265        // Provide a deterministic subnet principal for PocketIC runs.
266        let subnet_pid = Principal::from_slice(&[0xAA; 29]);
267        encode_one(SubnetIdentity::Manual(subnet_pid))
268            .map_err(|err| PublicError::internal(format!("encode_one failed: {err}")))
269    } else {
270        // Provide a minimal, deterministic env payload for standalone installs.
271        let root_pid = Principal::from_slice(&[0xBB; 29]);
272        let subnet_pid = Principal::from_slice(&[0xAA; 29]);
273        let env = EnvView {
274            prime_root_pid: Some(root_pid),
275            subnet_role: Some(SubnetRole::PRIME),
276            subnet_pid: Some(subnet_pid),
277            root_pid: Some(root_pid),
278            canister_role: Some(role),
279            parent_pid: Some(root_pid),
280        };
281
282        // Intentional: local standalone installs don't need directory views unless a test
283        // exercises directory-dependent auth/endpoints.
284        let payload = CanisterInitPayload {
285            env,
286            app_directory: AppDirectoryView(Vec::new()),
287            subnet_directory: SubnetDirectoryView(Vec::new()),
288        };
289        encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
290            .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))
291    }?;
292
293    Ok(args)
294}
295
296fn install_args_with_directories(
297    role: CanisterRole,
298    app_directory: AppDirectoryView,
299    subnet_directory: SubnetDirectoryView,
300) -> Result<Vec<u8>, PublicError> {
301    let args = if role.is_root() {
302        let subnet_pid = Principal::from_slice(&[0xAA; 29]);
303        encode_one(SubnetIdentity::Manual(subnet_pid))
304            .map_err(|err| PublicError::internal(format!("encode_one failed: {err}")))
305    } else {
306        let root_pid = Principal::from_slice(&[0xBB; 29]);
307        let subnet_pid = Principal::from_slice(&[0xAA; 29]);
308        let env = EnvView {
309            prime_root_pid: Some(root_pid),
310            subnet_role: Some(SubnetRole::PRIME),
311            subnet_pid: Some(subnet_pid),
312            root_pid: Some(root_pid),
313            canister_role: Some(role),
314            parent_pid: Some(root_pid),
315        };
316        let payload = CanisterInitPayload {
317            env,
318            app_directory,
319            subnet_directory,
320        };
321        encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
322            .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))
323    }?;
324
325    Ok(args)
326}