1use candid::{CandidType, decode_one, encode_args, encode_one, utils::ArgumentEncoder};
2use canic::{
3 Error,
4 core::{ids::CanisterRole, model::memory::topology::SubnetIdentity, ops::CanisterInitPayload},
5 types::Principal,
6};
7use derive_more::{Deref, DerefMut};
8use pocket_ic::{PocketIc, PocketIcBuilder};
9use serde::de::DeserializeOwned;
10
11pub struct PicBuilder(PocketIcBuilder);
16
17#[allow(clippy::new_without_default)]
18impl PicBuilder {
19 #[must_use]
21 pub fn new() -> Self {
22 Self(PocketIcBuilder::new())
23 }
24
25 #[must_use]
26 pub fn with_application_subnet(mut self) -> Self {
27 self.0 = self.0.with_application_subnet();
28 self
29 }
30
31 #[must_use]
32 pub fn with_nns_subnet(mut self) -> Self {
33 self.0 = self.0.with_nns_subnet();
34 self
35 }
36
37 #[must_use]
39 pub fn build(self) -> Pic {
40 Pic(self.0.build())
41 }
42}
43
44#[derive(Deref, DerefMut)]
49pub struct Pic(PocketIc);
50
51impl Pic {
52 pub fn create_and_install_canister(
54 &self,
55 ty: CanisterRole,
56 wasm: Vec<u8>,
57 ) -> Result<Principal, Error> {
58 let canister_id = self.create_canister();
60 self.add_cycles(canister_id, 1_000_000_000_000);
61
62 let init_bytes = install_args(ty)?;
64 self.0.install_canister(canister_id, wasm, init_bytes, None);
65
66 Ok(canister_id)
67 }
68
69 pub fn update_call<T, A>(
71 &self,
72 canister_id: Principal,
73 method: &str,
74 args: A,
75 ) -> Result<T, Error>
76 where
77 T: CandidType + DeserializeOwned,
78 A: ArgumentEncoder,
79 {
80 let bytes: Vec<u8> = encode_args(args)?;
81 let result = self
82 .0
83 .update_call(canister_id, Principal::anonymous(), method, bytes)
84 .map_err(|e| Error::test(e.to_string()))?;
85
86 decode_one(&result).map_err(Into::into)
87 }
88
89 pub fn query_call<T, A>(
91 &self,
92 canister_id: Principal,
93 method: &str,
94 args: A,
95 ) -> Result<T, Error>
96 where
97 T: CandidType + DeserializeOwned,
98 A: ArgumentEncoder,
99 {
100 let bytes: Vec<u8> = encode_args(args)?;
101 let result = self
102 .0
103 .query_call(canister_id, Principal::anonymous(), method, bytes)
104 .map_err(|e| Error::test(e.to_string()))?;
105
106 decode_one(&result).map_err(Into::into)
107 }
108}
109
110fn install_args(ty: CanisterRole) -> Result<Vec<u8>, Error> {
114 let args = if ty.is_root() {
115 let subnet_pid = Principal::from_slice(&[0xAA; 29]);
117 encode_one(SubnetIdentity::Manual(subnet_pid))
118 } else {
119 let payload = CanisterInitPayload::empty();
120 encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
121 }?;
122
123 Ok(args)
124}