canic_testkit/
pic.rs

1use candid::{CandidType, decode_one, encode_args, encode_one, utils::ArgumentEncoder};
2use canic::{
3    Error, core::types::Principal, model::memory::topology::SubnetIdentity,
4    ops::CanisterInitPayload, types::CanisterType,
5};
6use derive_more::{Deref, DerefMut};
7use pocket_ic::{PocketIc, PocketIcBuilder};
8use serde::de::DeserializeOwned;
9
10///
11/// PicBuilder
12///
13
14pub struct PicBuilder(PocketIcBuilder);
15
16#[allow(clippy::new_without_default)]
17impl PicBuilder {
18    /// Start a new PicBuilder with sensible defaults
19    #[must_use]
20    pub fn new() -> Self {
21        Self(PocketIcBuilder::new())
22    }
23
24    #[must_use]
25    pub fn with_application_subnet(mut self) -> Self {
26        self.0 = self.0.with_application_subnet();
27        self
28    }
29
30    #[must_use]
31    pub fn with_nns_subnet(mut self) -> Self {
32        self.0 = self.0.with_nns_subnet();
33        self
34    }
35
36    /// Finish building the PocketIC instance and wrap it
37    #[must_use]
38    pub fn build(self) -> Pic {
39        Pic(self.0.build())
40    }
41}
42
43///
44/// Pic
45///
46
47#[derive(Deref, DerefMut)]
48pub struct Pic(PocketIc);
49
50impl Pic {
51    /// Install a canister with the given type and wasm bytes
52    pub fn create_and_install_canister(
53        &self,
54        ty: CanisterType,
55        wasm: Vec<u8>,
56    ) -> Result<Principal, Error> {
57        // Create and fund the canister
58        let canister_id = self.create_canister();
59        self.add_cycles(canister_id, 1_000_000_000_000);
60
61        // Install
62        let init_bytes = install_args(ty)?;
63        self.0.install_canister(canister_id, wasm, init_bytes, None);
64
65        Ok(canister_id)
66    }
67
68    /// Generic update call helper (serializes args + decodes result)
69    pub fn update_call<T, A>(
70        &self,
71        canister_id: Principal,
72        method: &str,
73        args: A,
74    ) -> Result<T, Error>
75    where
76        T: CandidType + DeserializeOwned,
77        A: ArgumentEncoder,
78    {
79        let bytes: Vec<u8> = encode_args(args)?;
80        let result = self
81            .0
82            .update_call(canister_id, Principal::anonymous(), method, bytes)
83            .map_err(|e| Error::test(e.to_string()))?;
84
85        decode_one(&result).map_err(Into::into)
86    }
87
88    /// Generic query call helper
89    pub fn query_call<T, A>(
90        &self,
91        canister_id: Principal,
92        method: &str,
93        args: A,
94    ) -> Result<T, Error>
95    where
96        T: CandidType + DeserializeOwned,
97        A: ArgumentEncoder,
98    {
99        let bytes: Vec<u8> = encode_args(args)?;
100        let result = self
101            .0
102            .query_call(canister_id, Principal::anonymous(), method, bytes)
103            .map_err(|e| Error::test(e.to_string()))?;
104
105        decode_one(&result).map_err(Into::into)
106    }
107}
108
109/// --------------------------------------
110/// install_args helper
111/// --------------------------------------
112fn install_args(ty: CanisterType) -> Result<Vec<u8>, Error> {
113    let args = if ty.is_root() {
114        // Provide a deterministic subnet principal for PocketIC runs
115        let subnet_pid = Principal::from_slice(&[0xAA; 29]);
116        encode_one(SubnetIdentity::Manual(subnet_pid))
117    } else {
118        let payload = CanisterInitPayload::empty();
119        encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
120    }?;
121
122    Ok(args)
123}