canic_testkit/
pic.rs

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