1use candid::{CandidType, decode_one, encode_args, encode_one, utils::ArgumentEncoder};
2use canic::{
3 Error,
4 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
12pub struct PicBuilder(PocketIcBuilder);
17
18#[allow(clippy::new_without_default)]
19impl PicBuilder {
20 #[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 #[must_use]
40 pub fn build(self) -> Pic {
41 Pic(self.0.build())
42 }
43}
44
45#[derive(Deref, DerefMut)]
50pub struct Pic(PocketIc);
51
52impl Pic {
53 pub fn create_and_install_canister(
55 &self,
56 ty: CanisterType,
57 wasm: Vec<u8>,
58 ) -> Result<Principal, Error> {
59 let canister_id = self.create_canister();
61 self.add_cycles(canister_id, 1_000_000_000_000);
62
63 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 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 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
111fn 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}