canic_testkit/
pic.rs

1use candid::{CandidType, Principal, decode_one, encode_args, encode_one, utils::ArgumentEncoder};
2use canic::{
3    Error,
4    core::{
5        ids::{CanisterRole, SubnetRole},
6        ops::storage::{
7            CanisterInitPayload, directory::DirectoryView, env::EnvData, topology::SubnetIdentity,
8        },
9    },
10};
11use derive_more::{Deref, DerefMut};
12use pocket_ic::{PocketIc, PocketIcBuilder};
13use serde::de::DeserializeOwned;
14use std::sync::OnceLock;
15
16///
17/// PocketIC singleton
18///
19/// This crate models a *single* IC universe shared by all tests.
20/// We intentionally reuse one `PocketIc` instance to preserve determinism and
21/// to match the real IC's global, long-lived state.
22///
23/// Invariants:
24/// - Exactly one `PocketIc` instance exists for the entire test run.
25/// - All tests share the same universe (no resets between tests).
26/// - Tests are single-threaded and must not assume isolation.
27/// - Determinism is prioritized over per-test cleanliness.
28///
29/// The `OnceLock` is not about performance; it encodes these invariants so
30/// tests cannot accidentally spin up extra universes.
31///
32static PIC: OnceLock<Pic> = OnceLock::new();
33
34///
35/// Access the singleton PocketIC wrapper.
36///
37/// The global instance is created on first use and then reused.
38///
39#[must_use]
40pub fn pic() -> &'static Pic {
41    PIC.get_or_init(|| PicBuilder::new().with_application_subnet().build())
42}
43
44///
45/// PicBuilder
46/// Thin wrapper around the PocketIC builder.
47///
48/// This builder is only used to configure the singleton. It does not create
49/// additional IC instances beyond the global `Pic`.
50///
51/// Note: this file is test-only infrastructure; simplicity wins over abstraction.
52///
53pub struct PicBuilder(PocketIcBuilder);
54
55#[allow(clippy::new_without_default)]
56impl PicBuilder {
57    /// Start a new PicBuilder with sensible defaults.
58    #[must_use]
59    pub fn new() -> Self {
60        Self(PocketIcBuilder::new())
61    }
62
63    /// Include an application subnet in the singleton universe.
64    #[must_use]
65    pub fn with_application_subnet(mut self) -> Self {
66        self.0 = self.0.with_application_subnet();
67        self
68    }
69
70    /// Include an NNS subnet in the singleton universe.
71    #[must_use]
72    pub fn with_nns_subnet(mut self) -> Self {
73        self.0 = self.0.with_nns_subnet();
74        self
75    }
76
77    /// Finish building the singleton PocketIC instance and wrap it.
78    #[must_use]
79    pub fn build(self) -> Pic {
80        Pic(self.0.build())
81    }
82}
83
84///
85/// Pic
86/// Thin wrapper around the global PocketIC instance.
87///
88/// This type intentionally exposes only a minimal API surface; callers should
89/// use `pic()` to obtain the singleton and then perform installs/calls.
90///
91#[derive(Deref, DerefMut)]
92pub struct Pic(PocketIc);
93
94impl Pic {
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, Error> {
103        // Create and fund the canister.
104        let canister_id = self.create_canister();
105        self.add_cycles(canister_id, 1_000_000_000_000);
106
107        // Install with deterministic init arguments.
108        let init_bytes = install_args(role)?;
109        self.0.install_canister(canister_id, wasm, init_bytes, None);
110
111        Ok(canister_id)
112    }
113
114    /// Install a canister with a custom directory snapshot (local-only helper).
115    ///
116    /// Use this when a test exercises directory-dependent auth/endpoints and
117    /// cannot rely on root to provide a snapshot.
118    pub fn create_and_install_canister_with_directories(
119        &self,
120        role: CanisterRole,
121        wasm: Vec<u8>,
122        app_directory: DirectoryView,
123        subnet_directory: DirectoryView,
124    ) -> Result<Principal, Error> {
125        let canister_id = self.create_canister();
126        self.add_cycles(canister_id, 1_000_000_000_000);
127
128        let init_bytes = install_args_with_directories(role, app_directory, subnet_directory)?;
129        self.0.install_canister(canister_id, wasm, init_bytes, None);
130
131        Ok(canister_id)
132    }
133
134    /// Generic update call helper (serializes args + decodes result).
135    pub fn update_call<T, A>(
136        &self,
137        canister_id: Principal,
138        method: &str,
139        args: A,
140    ) -> Result<T, Error>
141    where
142        T: CandidType + DeserializeOwned,
143        A: ArgumentEncoder,
144    {
145        let bytes: Vec<u8> = encode_args(args)?;
146        let result = self
147            .0
148            .update_call(canister_id, Principal::anonymous(), method, bytes)
149            .map_err(|e| Error::test(e.to_string()))?;
150
151        decode_one(&result).map_err(Into::into)
152    }
153
154    /// Generic query call helper.
155    pub fn query_call<T, A>(
156        &self,
157        canister_id: Principal,
158        method: &str,
159        args: A,
160    ) -> Result<T, Error>
161    where
162        T: CandidType + DeserializeOwned,
163        A: ArgumentEncoder,
164    {
165        let bytes: Vec<u8> = encode_args(args)?;
166        let result = self
167            .0
168            .query_call(canister_id, Principal::anonymous(), method, bytes)
169            .map_err(|e| Error::test(e.to_string()))?;
170
171        decode_one(&result).map_err(Into::into)
172    }
173}
174
175/// --------------------------------------
176/// install_args helper
177/// --------------------------------------
178///
179/// Init semantics:
180/// - Root canisters receive a `SubnetIdentity` (direct root bootstrap).
181/// - Non-root canisters receive `EnvData` + optional directory snapshots.
182///
183/// Directory handling:
184/// - By default, directory views are empty for standalone installs.
185/// - Directory-dependent logic is opt-in via `install_args_with_directories`.
186/// - Root-provisioned installs will populate directories via cascade.
187///
188fn install_args(role: CanisterRole) -> Result<Vec<u8>, Error> {
189    let args = if role.is_root() {
190        // Provide a deterministic subnet principal for PocketIC runs.
191        let subnet_pid = Principal::from_slice(&[0xAA; 29]);
192        encode_one(SubnetIdentity::Manual(subnet_pid))
193    } else {
194        // Provide a minimal, deterministic env payload for standalone installs.
195        let root_pid = Principal::from_slice(&[0xBB; 29]);
196        let subnet_pid = Principal::from_slice(&[0xAA; 29]);
197        let env = EnvData {
198            prime_root_pid: Some(root_pid),
199            subnet_role: Some(SubnetRole::PRIME),
200            subnet_pid: Some(subnet_pid),
201            root_pid: Some(root_pid),
202            canister_role: Some(role),
203            parent_pid: Some(root_pid),
204        };
205        // Intentional: local standalone installs don't need directory views unless a test
206        // exercises directory-dependent auth/endpoints.
207        let payload = CanisterInitPayload::new(env, Vec::new(), Vec::new());
208        encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
209    }?;
210
211    Ok(args)
212}
213
214fn install_args_with_directories(
215    role: CanisterRole,
216    app_directory: DirectoryView,
217    subnet_directory: DirectoryView,
218) -> Result<Vec<u8>, Error> {
219    let args = if role.is_root() {
220        let subnet_pid = Principal::from_slice(&[0xAA; 29]);
221        encode_one(SubnetIdentity::Manual(subnet_pid))
222    } else {
223        let root_pid = Principal::from_slice(&[0xBB; 29]);
224        let subnet_pid = Principal::from_slice(&[0xAA; 29]);
225        let env = EnvData {
226            prime_root_pid: Some(root_pid),
227            subnet_role: Some(SubnetRole::PRIME),
228            subnet_pid: Some(subnet_pid),
229            root_pid: Some(root_pid),
230            canister_role: Some(role),
231            parent_pid: Some(root_pid),
232        };
233        let payload = CanisterInitPayload::new(env, app_directory, subnet_directory);
234        encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
235    }?;
236
237    Ok(args)
238}