Skip to main content

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::EnvBootstrapArgs,
8        subnet::SubnetIdentity,
9        topology::{AppDirectoryArgs, SubnetDirectoryArgs},
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    fn create_funded_and_install(&self, wasm: Vec<u8>, init_bytes: Vec<u8>) -> Principal {
104        let canister_id = self.create_canister();
105        self.add_cycles(canister_id, INSTALL_CYCLES);
106        self.0.install_canister(canister_id, wasm, init_bytes, None);
107
108        canister_id
109    }
110
111    /// Generic update call helper (serializes args + decodes result).
112    pub fn update_call<T, A>(
113        &self,
114        canister_id: Principal,
115        method: &str,
116        args: A,
117    ) -> Result<T, Error>
118    where
119        T: CandidType + DeserializeOwned,
120        A: ArgumentEncoder,
121    {
122        let bytes: Vec<u8> = encode_args(args)
123            .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
124        let result = self
125            .0
126            .update_call(canister_id, Principal::anonymous(), method, bytes)
127            .map_err(|err| {
128                Error::internal(format!(
129                    "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
130                ))
131            })?;
132
133        decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
134    }
135
136    /// Generic update call helper with an explicit caller principal.
137    pub fn update_call_as<T, A>(
138        &self,
139        canister_id: Principal,
140        caller: Principal,
141        method: &str,
142        args: A,
143    ) -> Result<T, Error>
144    where
145        T: CandidType + DeserializeOwned,
146        A: ArgumentEncoder,
147    {
148        let bytes: Vec<u8> = encode_args(args)
149            .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
150        let result = self
151            .0
152            .update_call(canister_id, caller, method, bytes)
153            .map_err(|err| {
154                Error::internal(format!(
155                    "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
156                ))
157            })?;
158
159        decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
160    }
161
162    /// Generic query call helper.
163    pub fn query_call<T, A>(
164        &self,
165        canister_id: Principal,
166        method: &str,
167        args: A,
168    ) -> Result<T, Error>
169    where
170        T: CandidType + DeserializeOwned,
171        A: ArgumentEncoder,
172    {
173        let bytes: Vec<u8> = encode_args(args)
174            .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
175        let result = self
176            .0
177            .query_call(canister_id, Principal::anonymous(), method, bytes)
178            .map_err(|err| {
179                Error::internal(format!(
180                    "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
181                ))
182            })?;
183
184        decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
185    }
186
187    /// Generic query call helper with an explicit caller principal.
188    pub fn query_call_as<T, A>(
189        &self,
190        canister_id: Principal,
191        caller: Principal,
192        method: &str,
193        args: A,
194    ) -> Result<T, Error>
195    where
196        T: CandidType + DeserializeOwned,
197        A: ArgumentEncoder,
198    {
199        let bytes: Vec<u8> = encode_args(args)
200            .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
201        let result = self
202            .0
203            .query_call(canister_id, caller, method, bytes)
204            .map_err(|err| {
205                Error::internal(format!(
206                    "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
207                ))
208            })?;
209
210        decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
211    }
212
213    pub fn tick_n(&self, times: usize) {
214        for _ in 0..times {
215            self.tick();
216        }
217    }
218}
219
220/// --------------------------------------
221/// install_args helper
222/// --------------------------------------
223///
224/// Init semantics:
225/// - Root canisters receive a `SubnetIdentity` (direct root bootstrap).
226/// - Non-root canisters receive `EnvBootstrapArgs` + optional directory snapshots.
227///
228/// Directory handling:
229/// - By default, directory views are empty for standalone installs.
230/// - Directory-dependent logic is opt-in via `install_args_with_directories`.
231/// - Root-provisioned installs will populate directories via cascade.
232///
233fn install_args(role: CanisterRole) -> Result<Vec<u8>, Error> {
234    if role.is_root() {
235        install_root_args()
236    } else {
237        // Non-root standalone install.
238        // Provide only what is structurally known at install time.
239        let env = EnvBootstrapArgs {
240            prime_root_pid: None,
241            subnet_role: None,
242            subnet_pid: None,
243            root_pid: None,
244            canister_role: Some(role),
245            parent_pid: None,
246        };
247
248        // Intentional: standalone installs do not require directories unless
249        // a test explicitly exercises directory-dependent behavior.
250        let payload = CanisterInitPayload {
251            env,
252            app_directory: AppDirectoryArgs(Vec::new()),
253            subnet_directory: SubnetDirectoryArgs(Vec::new()),
254        };
255
256        encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
257            .map_err(|err| Error::internal(format!("encode_args failed: {err}")))
258    }
259}
260
261fn install_root_args() -> Result<Vec<u8>, Error> {
262    encode_one(SubnetIdentity::Manual)
263        .map_err(|err| Error::internal(format!("encode_one failed: {err}")))
264}