use candid::{Principal, encode_args, encode_one};
use canic::{
Error,
cdk::types::TC,
dto::{
abi::v1::CanisterInitPayload,
env::EnvBootstrapArgs,
subnet::SubnetIdentity,
topology::{AppIndexArgs, SubnetIndexArgs},
},
ids::CanisterRole,
};
use pocket_ic::{
CanisterStatusResult, PocketIc, PocketIcBuilder, RejectResponse, common::rest::RawMessageId,
};
use std::time::Duration;
mod baseline;
mod calls;
mod diagnostics;
mod errors;
mod lifecycle;
mod process_lock;
mod readiness;
mod snapshot;
mod standalone;
mod startup;
pub use baseline::{
CachedPicBaseline, CachedPicBaselineGuard, ControllerSnapshots,
restore_or_rebuild_cached_pic_baseline,
};
pub use errors::{PicInstallError, StandaloneCanisterFixtureError};
pub use process_lock::{
PicSerialGuard, PicSerialGuardError, acquire_pic_serial_guard, try_acquire_pic_serial_guard,
};
pub use readiness::{role_pid, wait_until_ready};
pub use startup::PicStartError;
const INSTALL_CYCLES: u128 = 500 * TC;
pub use standalone::{
StandaloneCanisterFixture, install_prebuilt_canister, install_prebuilt_canister_with_cycles,
install_standalone_canister, try_install_prebuilt_canister,
try_install_prebuilt_canister_with_cycles,
};
#[must_use]
pub fn pic() -> Pic {
try_pic().unwrap_or_else(|err| panic!("failed to start PocketIC: {err}"))
}
pub fn try_pic() -> Result<Pic, PicStartError> {
PicBuilder::new().with_application_subnet().try_build()
}
pub struct PicBuilder(PocketIcBuilder);
#[expect(clippy::new_without_default)]
impl PicBuilder {
#[must_use]
pub fn new() -> Self {
Self(PocketIcBuilder::new())
}
#[must_use]
pub fn with_application_subnet(mut self) -> Self {
self.0 = self.0.with_application_subnet();
self
}
#[must_use]
pub fn with_ii_subnet(mut self) -> Self {
self.0 = self.0.with_ii_subnet();
self
}
#[must_use]
pub fn with_nns_subnet(mut self) -> Self {
self.0 = self.0.with_nns_subnet();
self
}
#[must_use]
pub fn build(self) -> Pic {
self.try_build()
.unwrap_or_else(|err| panic!("failed to start PocketIC: {err}"))
}
pub fn try_build(self) -> Result<Pic, PicStartError> {
startup::try_build_pic(self.0)
}
}
pub struct Pic {
inner: PocketIc,
}
impl Pic {
pub fn tick(&self) {
self.inner.tick();
}
pub fn advance_time(&self, duration: Duration) {
self.inner.advance_time(duration);
}
#[must_use]
pub fn create_canister(&self) -> Principal {
self.inner.create_canister()
}
pub fn add_cycles(&self, canister_id: Principal, amount: u128) {
let _ = self.inner.add_cycles(canister_id, amount);
}
pub fn install_canister(
&self,
canister_id: Principal,
wasm_module: Vec<u8>,
arg: Vec<u8>,
sender: Option<Principal>,
) {
self.inner
.install_canister(canister_id, wasm_module, arg, sender);
}
pub fn upgrade_canister(
&self,
canister_id: Principal,
wasm_module: Vec<u8>,
arg: Vec<u8>,
sender: Option<Principal>,
) -> Result<(), RejectResponse> {
self.inner
.upgrade_canister(canister_id, wasm_module, arg, sender)
}
pub fn reinstall_canister(
&self,
canister_id: Principal,
wasm_module: Vec<u8>,
arg: Vec<u8>,
sender: Option<Principal>,
) -> Result<(), RejectResponse> {
self.inner
.reinstall_canister(canister_id, wasm_module, arg, sender)
}
pub fn submit_call(
&self,
canister_id: Principal,
sender: Principal,
method: &str,
payload: Vec<u8>,
) -> Result<RawMessageId, RejectResponse> {
self.inner.submit_call(canister_id, sender, method, payload)
}
pub fn await_call(&self, message_id: RawMessageId) -> Result<Vec<u8>, RejectResponse> {
self.inner.await_call(message_id)
}
pub fn canister_status(
&self,
canister_id: Principal,
sender: Option<Principal>,
) -> Result<CanisterStatusResult, RejectResponse> {
self.inner.canister_status(canister_id, sender)
}
pub fn fetch_canister_logs(
&self,
canister_id: Principal,
sender: Principal,
) -> Result<Vec<pocket_ic::CanisterLogRecord>, RejectResponse> {
self.inner.fetch_canister_logs(canister_id, sender)
}
#[must_use]
pub fn current_time_nanos(&self) -> u64 {
self.inner.get_time().as_nanos_since_unix_epoch()
}
pub fn restore_time_nanos(&self, nanos_since_epoch: u64) {
let restored = pocket_ic::Time::from_nanos_since_unix_epoch(nanos_since_epoch);
self.inner.set_time(restored);
self.inner.set_certified_time(restored);
}
}
fn install_args(role: CanisterRole) -> Result<Vec<u8>, Error> {
if role.is_root() {
install_root_args()
} else {
let env = EnvBootstrapArgs {
prime_root_pid: None,
subnet_role: None,
subnet_pid: None,
root_pid: None,
canister_role: Some(role),
parent_pid: None,
};
let payload = CanisterInitPayload {
env,
app_index: AppIndexArgs(Vec::new()),
subnet_index: SubnetIndexArgs(Vec::new()),
};
encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
.map_err(|err| Error::internal(format!("encode_args failed: {err}")))
}
}
fn install_root_args() -> Result<Vec<u8>, Error> {
encode_one(SubnetIdentity::Manual)
.map_err(|err| Error::internal(format!("encode_one failed: {err}")))
}