Skip to main content

canic_testkit/pic/
mod.rs

1use candid::{Principal, decode_one, encode_args, encode_one};
2use canic::{
3    Error,
4    cdk::types::TC,
5    dto::{
6        abi::v1::CanisterInitPayload,
7        env::EnvBootstrapArgs,
8        subnet::SubnetIdentity,
9        topology::{AppDirectoryArgs, SubnetDirectoryArgs, SubnetRegistryResponse},
10    },
11    ids::CanisterRole,
12    protocol,
13};
14use pocket_ic::{PocketIc, PocketIcBuilder};
15use std::{
16    ops::{Deref, DerefMut},
17    panic::AssertUnwindSafe,
18};
19
20mod baseline;
21mod calls;
22mod diagnostics;
23mod lifecycle;
24mod process_lock;
25mod snapshot;
26mod standalone;
27mod startup;
28
29pub use baseline::{
30    CachedPicBaseline, CachedPicBaselineGuard, ControllerSnapshots, acquire_cached_pic_baseline,
31    drop_stale_cached_pic_baseline, restore_or_rebuild_cached_pic_baseline,
32};
33pub use process_lock::{
34    PicSerialGuard, PicSerialGuardError, acquire_pic_serial_guard, try_acquire_pic_serial_guard,
35};
36pub use startup::PicStartError;
37const INSTALL_CYCLES: u128 = 500 * TC;
38
39///
40/// PicInstallError
41///
42
43#[derive(Debug, Eq, PartialEq)]
44pub struct PicInstallError {
45    canister_id: Principal,
46    message: String,
47}
48
49///
50/// StandaloneCanisterFixtureError
51///
52
53#[derive(Debug)]
54pub enum StandaloneCanisterFixtureError {
55    SerialGuard(PicSerialGuardError),
56    Start(PicStartError),
57    Install(PicInstallError),
58}
59
60pub use standalone::{
61    StandaloneCanisterFixture, install_prebuilt_canister, install_prebuilt_canister_with_cycles,
62    install_standalone_canister, try_install_prebuilt_canister,
63    try_install_prebuilt_canister_with_cycles,
64};
65
66///
67/// Create a fresh PocketIC universe.
68///
69/// IMPORTANT:
70/// - Each call creates a new IC instance
71/// - WARNING: callers must hold a `PicSerialGuard` for the full `Pic` lifetime
72/// - Required to avoid PocketIC wasm chunk store exhaustion
73///
74#[must_use]
75pub fn pic() -> Pic {
76    try_pic().unwrap_or_else(|err| panic!("failed to start PocketIC: {err}"))
77}
78
79/// Create a fresh PocketIC universe without panicking on startup failures.
80pub fn try_pic() -> Result<Pic, PicStartError> {
81    PicBuilder::new().with_application_subnet().try_build()
82}
83
84/// Wait until a PocketIC canister reports `canic_ready`.
85pub fn wait_until_ready(pic: &PocketIc, canister_id: Principal, tick_limit: usize) {
86    let payload = encode_args(()).expect("encode empty args");
87
88    for _ in 0..tick_limit {
89        if let Ok(bytes) = pic.query_call(
90            canister_id,
91            Principal::anonymous(),
92            protocol::CANIC_READY,
93            payload.clone(),
94        ) && let Ok(ready) = decode_one::<bool>(&bytes)
95            && ready
96        {
97            return;
98        }
99        pic.tick();
100    }
101
102    panic!("canister did not report ready in time: {canister_id}");
103}
104
105/// Resolve one role principal from root's subnet registry, polling until present.
106#[must_use]
107pub fn role_pid(
108    pic: &PocketIc,
109    root_id: Principal,
110    role: &'static str,
111    tick_limit: usize,
112) -> Principal {
113    for _ in 0..tick_limit {
114        let registry: Result<Result<SubnetRegistryResponse, Error>, Error> = {
115            let payload = encode_args(()).expect("encode empty args");
116            pic.query_call(
117                root_id,
118                Principal::anonymous(),
119                protocol::CANIC_SUBNET_REGISTRY,
120                payload,
121            )
122            .map_err(|err| {
123                Error::internal(format!(
124                    "pocket_ic query_call failed (canister={root_id}, method={}): {err}",
125                    protocol::CANIC_SUBNET_REGISTRY
126                ))
127            })
128            .and_then(|bytes| {
129                decode_one(&bytes).map_err(|err| {
130                    Error::internal(format!("decode_one failed for subnet registry: {err}"))
131                })
132            })
133        };
134
135        if let Ok(Ok(registry)) = registry
136            && let Some(pid) = registry
137                .0
138                .into_iter()
139                .find(|entry| entry.role == CanisterRole::new(role))
140                .map(|entry| entry.pid)
141        {
142            return pid;
143        }
144
145        pic.tick();
146    }
147
148    panic!("{role} canister must be registered");
149}
150
151///
152/// PicBuilder
153/// Thin wrapper around the PocketIC builder.
154///
155/// This builder is only used to configure the singleton. It does not create
156/// additional IC instances beyond the global `Pic`.
157///
158/// Note: this file is test-only infrastructure; simplicity wins over abstraction.
159///
160
161pub struct PicBuilder(PocketIcBuilder);
162
163#[expect(clippy::new_without_default)]
164impl PicBuilder {
165    /// Start a new PicBuilder with sensible defaults.
166    #[must_use]
167    pub fn new() -> Self {
168        Self(PocketIcBuilder::new())
169    }
170
171    /// Include an application subnet in the PocketIC universe.
172    #[must_use]
173    pub fn with_application_subnet(mut self) -> Self {
174        self.0 = self.0.with_application_subnet();
175        self
176    }
177
178    /// Include an II subnet so threshold keys are available in the PocketIC universe.
179    #[must_use]
180    pub fn with_ii_subnet(mut self) -> Self {
181        self.0 = self.0.with_ii_subnet();
182        self
183    }
184
185    /// Include an NNS subnet in the PocketIC universe.
186    #[must_use]
187    pub fn with_nns_subnet(mut self) -> Self {
188        self.0 = self.0.with_nns_subnet();
189        self
190    }
191
192    /// Finish building the PocketIC instance and wrap it.
193    #[must_use]
194    pub fn build(self) -> Pic {
195        self.try_build()
196            .unwrap_or_else(|err| panic!("failed to start PocketIC: {err}"))
197    }
198
199    /// Finish building the PocketIC instance without panicking on startup failures.
200    pub fn try_build(self) -> Result<Pic, PicStartError> {
201        startup::try_build_pic(AssertUnwindSafe(self.0).0)
202    }
203}
204
205impl PicInstallError {
206    /// Capture one install failure for a specific canister id.
207    #[must_use]
208    pub const fn new(canister_id: Principal, message: String) -> Self {
209        Self {
210            canister_id,
211            message,
212        }
213    }
214
215    /// Read the canister id that failed to install.
216    #[must_use]
217    pub const fn canister_id(&self) -> Principal {
218        self.canister_id
219    }
220
221    /// Read the captured panic message from the install attempt.
222    #[must_use]
223    pub fn message(&self) -> &str {
224        &self.message
225    }
226}
227
228impl std::fmt::Display for PicInstallError {
229    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
230        write!(
231            f,
232            "failed to install canister {}: {}",
233            self.canister_id, self.message
234        )
235    }
236}
237
238impl std::error::Error for PicInstallError {}
239
240impl std::fmt::Display for StandaloneCanisterFixtureError {
241    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
242        match self {
243            Self::SerialGuard(err) => write!(f, "{err}"),
244            Self::Start(err) => write!(f, "{err}"),
245            Self::Install(err) => write!(f, "{err}"),
246        }
247    }
248}
249
250impl std::error::Error for StandaloneCanisterFixtureError {
251    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
252        match self {
253            Self::SerialGuard(err) => Some(err),
254            Self::Start(err) => Some(err),
255            Self::Install(err) => Some(err),
256        }
257    }
258}
259
260///
261/// Pic
262/// Thin wrapper around a PocketIC instance.
263///
264/// This type intentionally exposes only a minimal API surface; callers should
265/// use `pic()` to obtain an instance and then perform installs/calls.
266/// Callers must hold a `PicSerialGuard` for the full `Pic` lifetime.
267///
268
269pub struct Pic {
270    inner: PocketIc,
271}
272
273impl Pic {
274    /// Capture the current PocketIC wall-clock time as nanoseconds since epoch.
275    #[must_use]
276    pub fn current_time_nanos(&self) -> u64 {
277        self.inner.get_time().as_nanos_since_unix_epoch()
278    }
279
280    /// Restore PocketIC wall-clock and certified time from a captured nanosecond value.
281    pub fn restore_time_nanos(&self, nanos_since_epoch: u64) {
282        let restored = pocket_ic::Time::from_nanos_since_unix_epoch(nanos_since_epoch);
283        self.inner.set_time(restored);
284        self.inner.set_certified_time(restored);
285    }
286}
287
288impl Deref for Pic {
289    type Target = PocketIc;
290
291    fn deref(&self) -> &Self::Target {
292        &self.inner
293    }
294}
295
296impl DerefMut for Pic {
297    fn deref_mut(&mut self) -> &mut Self::Target {
298        &mut self.inner
299    }
300}
301
302/// --------------------------------------
303/// install_args helper
304/// --------------------------------------
305///
306/// Init semantics:
307/// - Root canisters receive a `SubnetIdentity` (direct root bootstrap).
308/// - Non-root canisters receive `EnvBootstrapArgs` + optional directory snapshots.
309///
310/// Directory handling:
311/// - By default, directory views are empty for standalone installs.
312/// - Directory-dependent logic is opt-in via `install_args_with_directories`.
313/// - Root-provisioned installs will populate directories via cascade.
314///
315
316fn install_args(role: CanisterRole) -> Result<Vec<u8>, Error> {
317    if role.is_root() {
318        install_root_args()
319    } else {
320        // Non-root standalone install.
321        // Provide only what is structurally known at install time.
322        let env = EnvBootstrapArgs {
323            prime_root_pid: None,
324            subnet_role: None,
325            subnet_pid: None,
326            root_pid: None,
327            canister_role: Some(role),
328            parent_pid: None,
329        };
330
331        // Intentional: standalone installs do not require directories unless
332        // a test explicitly exercises directory-dependent behavior.
333        let payload = CanisterInitPayload {
334            env,
335            app_directory: AppDirectoryArgs(Vec::new()),
336            subnet_directory: SubnetDirectoryArgs(Vec::new()),
337        };
338
339        encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
340            .map_err(|err| Error::internal(format!("encode_args failed: {err}")))
341    }
342}
343
344fn install_root_args() -> Result<Vec<u8>, Error> {
345    encode_one(SubnetIdentity::Manual)
346        .map_err(|err| Error::internal(format!("encode_one failed: {err}")))
347}
348
349// Prefer the likely controller sender first to reduce noisy management-call failures.