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
10pub struct PicBuilder(PocketIcBuilder);
15
16#[allow(clippy::new_without_default)]
17impl PicBuilder {
18 #[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 #[must_use]
38 pub fn build(self) -> Pic {
39 Pic(self.0.build())
40 }
41}
42
43#[derive(Deref, DerefMut)]
48pub struct Pic(PocketIc);
49
50impl Pic {
51 pub fn create_and_install_canister(
53 &self,
54 ty: CanisterType,
55 wasm: Vec<u8>,
56 ) -> Result<Principal, Error> {
57 let canister_id = self.create_canister();
59 self.add_cycles(canister_id, 1_000_000_000_000);
60
61 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 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 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
109fn install_args(ty: CanisterType) -> Result<Vec<u8>, Error> {
113 let args = if ty.is_root() {
114 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}