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