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#[derive(Debug, Eq, PartialEq)]
44pub struct PicInstallError {
45 canister_id: Principal,
46 message: String,
47}
48
49#[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#[must_use]
75pub fn pic() -> Pic {
76 try_pic().unwrap_or_else(|err| panic!("failed to start PocketIC: {err}"))
77}
78
79pub fn try_pic() -> Result<Pic, PicStartError> {
81 PicBuilder::new().with_application_subnet().try_build()
82}
83
84pub 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#[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
151pub struct PicBuilder(PocketIcBuilder);
162
163#[expect(clippy::new_without_default)]
164impl PicBuilder {
165 #[must_use]
167 pub fn new() -> Self {
168 Self(PocketIcBuilder::new())
169 }
170
171 #[must_use]
173 pub fn with_application_subnet(mut self) -> Self {
174 self.0 = self.0.with_application_subnet();
175 self
176 }
177
178 #[must_use]
180 pub fn with_ii_subnet(mut self) -> Self {
181 self.0 = self.0.with_ii_subnet();
182 self
183 }
184
185 #[must_use]
187 pub fn with_nns_subnet(mut self) -> Self {
188 self.0 = self.0.with_nns_subnet();
189 self
190 }
191
192 #[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 pub fn try_build(self) -> Result<Pic, PicStartError> {
201 startup::try_build_pic(AssertUnwindSafe(self.0).0)
202 }
203}
204
205impl PicInstallError {
206 #[must_use]
208 pub const fn new(canister_id: Principal, message: String) -> Self {
209 Self {
210 canister_id,
211 message,
212 }
213 }
214
215 #[must_use]
217 pub const fn canister_id(&self) -> Principal {
218 self.canister_id
219 }
220
221 #[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
260pub struct Pic {
270 inner: PocketIc,
271}
272
273impl Pic {
274 #[must_use]
276 pub fn current_time_nanos(&self) -> u64 {
277 self.inner.get_time().as_nanos_since_unix_epoch()
278 }
279
280 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
302fn install_args(role: CanisterRole) -> Result<Vec<u8>, Error> {
317 if role.is_root() {
318 install_root_args()
319 } else {
320 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 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