canic_testkit/
pic.rs

1use candid::{CandidType, Principal, decode_one, encode_args, encode_one, utils::ArgumentEncoder};
2use canic::{
3    Error,
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(&self, wasm: Vec<u8>) -> Result<Principal, Error> {
85        let init_bytes = install_root_args()?;
86
87        Ok(self.create_funded_and_install(wasm, init_bytes))
88    }
89
90    /// Install a canister with the given type and wasm bytes.
91    ///
92    /// Install failures are treated as fatal in tests.
93    pub fn create_and_install_canister(
94        &self,
95        role: CanisterRole,
96        wasm: Vec<u8>,
97    ) -> Result<Principal, Error> {
98        let init_bytes = install_args(role)?;
99
100        Ok(self.create_funded_and_install(wasm, init_bytes))
101    }
102
103    /// Install a canister with a custom directory snapshot (local-only helper).
104    ///
105    /// Use this when a test exercises directory-dependent auth/endpoints and
106    /// cannot rely on root to provide a snapshot.
107    pub fn create_and_install_canister_with_directories(
108        &self,
109        role: CanisterRole,
110        wasm: Vec<u8>,
111        app_directory: AppDirectoryView,
112        subnet_directory: SubnetDirectoryView,
113    ) -> Result<Principal, Error> {
114        let init_bytes = install_args_with_directories(role, app_directory, subnet_directory)?;
115
116        Ok(self.create_funded_and_install(wasm, init_bytes))
117    }
118
119    fn create_funded_and_install(&self, wasm: Vec<u8>, init_bytes: Vec<u8>) -> Principal {
120        let canister_id = self.create_canister();
121        self.add_cycles(canister_id, INSTALL_CYCLES);
122        self.0.install_canister(canister_id, wasm, init_bytes, None);
123
124        canister_id
125    }
126
127    /// Generic update call helper (serializes args + decodes result).
128    pub fn update_call<T, A>(
129        &self,
130        canister_id: Principal,
131        method: &str,
132        args: A,
133    ) -> Result<T, Error>
134    where
135        T: CandidType + DeserializeOwned,
136        A: ArgumentEncoder,
137    {
138        let bytes: Vec<u8> = encode_args(args)
139            .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
140        let result = self
141            .0
142            .update_call(canister_id, Principal::anonymous(), method, bytes)
143            .map_err(|err| {
144                Error::internal(format!(
145                    "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
146                ))
147            })?;
148
149        decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
150    }
151
152    /// Generic update call helper with an explicit caller principal.
153    pub fn update_call_as<T, A>(
154        &self,
155        canister_id: Principal,
156        caller: Principal,
157        method: &str,
158        args: A,
159    ) -> Result<T, Error>
160    where
161        T: CandidType + DeserializeOwned,
162        A: ArgumentEncoder,
163    {
164        let bytes: Vec<u8> = encode_args(args)
165            .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
166        let result = self
167            .0
168            .update_call(canister_id, caller, method, bytes)
169            .map_err(|err| {
170                Error::internal(format!(
171                    "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
172                ))
173            })?;
174
175        decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
176    }
177
178    /// Generic query call helper.
179    pub fn query_call<T, A>(
180        &self,
181        canister_id: Principal,
182        method: &str,
183        args: A,
184    ) -> Result<T, Error>
185    where
186        T: CandidType + DeserializeOwned,
187        A: ArgumentEncoder,
188    {
189        let bytes: Vec<u8> = encode_args(args)
190            .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
191        let result = self
192            .0
193            .query_call(canister_id, Principal::anonymous(), method, bytes)
194            .map_err(|err| {
195                Error::internal(format!(
196                    "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
197                ))
198            })?;
199
200        decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
201    }
202
203    /// Generic query call helper with an explicit caller principal.
204    pub fn query_call_as<T, A>(
205        &self,
206        canister_id: Principal,
207        caller: Principal,
208        method: &str,
209        args: A,
210    ) -> Result<T, Error>
211    where
212        T: CandidType + DeserializeOwned,
213        A: ArgumentEncoder,
214    {
215        let bytes: Vec<u8> = encode_args(args)
216            .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
217        let result = self
218            .0
219            .query_call(canister_id, caller, method, bytes)
220            .map_err(|err| {
221                Error::internal(format!(
222                    "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
223                ))
224            })?;
225
226        decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
227    }
228
229    pub fn tick_n(&self, times: usize) {
230        for _ in 0..times {
231            self.tick();
232        }
233    }
234}
235
236/// --------------------------------------
237/// install_args helper
238/// --------------------------------------
239///
240/// Init semantics:
241/// - Root canisters receive a `SubnetIdentity` (direct root bootstrap).
242/// - Non-root canisters receive `EnvView` + optional directory snapshots.
243///
244/// Directory handling:
245/// - By default, directory views are empty for standalone installs.
246/// - Directory-dependent logic is opt-in via `install_args_with_directories`.
247/// - Root-provisioned installs will populate directories via cascade.
248///
249fn install_args(role: CanisterRole) -> Result<Vec<u8>, Error> {
250    if role.is_root() {
251        install_root_args()
252    } else {
253        // Non-root standalone install.
254        // Provide only what is structurally known at install time.
255        let env = EnvView {
256            prime_root_pid: None,
257            subnet_role: None,
258            subnet_pid: None,
259            root_pid: None,
260            canister_role: Some(role),
261            parent_pid: None,
262        };
263
264        // Intentional: standalone installs do not require directories unless
265        // a test explicitly exercises directory-dependent behavior.
266        let payload = CanisterInitPayload {
267            env,
268            app_directory: AppDirectoryView(Vec::new()),
269            subnet_directory: SubnetDirectoryView(Vec::new()),
270        };
271
272        encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
273            .map_err(|err| Error::internal(format!("encode_args failed: {err}")))
274    }
275}
276
277fn install_root_args() -> Result<Vec<u8>, Error> {
278    encode_one(SubnetIdentity::Manual)
279        .map_err(|err| Error::internal(format!("encode_one failed: {err}")))
280}
281
282fn install_args_with_directories(
283    role: CanisterRole,
284    app_directory: AppDirectoryView,
285    subnet_directory: SubnetDirectoryView,
286) -> Result<Vec<u8>, Error> {
287    if role.is_root() {
288        // Root canister: runtime identity only.
289        // No fake principals. Runtime/bootstrap will resolve actual context.
290        encode_one(SubnetIdentity::Manual)
291            .map_err(|err| Error::internal(format!("encode_one failed: {err}")))
292    } else {
293        // Non-root canister: pass structural context, not invented identities.
294        let env = EnvView {
295            prime_root_pid: None,
296            subnet_role: None,
297            subnet_pid: None,
298            root_pid: None,
299            canister_role: Some(role),
300            parent_pid: None,
301        };
302
303        let payload = CanisterInitPayload {
304            env,
305            app_directory,
306            subnet_directory,
307        };
308
309        encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
310            .map_err(|err| Error::internal(format!("encode_args failed: {err}")))
311    }
312}