canic_testkit/
pic.rs

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