1use candid::{CandidType, Principal, decode_one, encode_args, encode_one, utils::ArgumentEncoder};
2use canic::{
3 Error,
4 cdk::types::TC,
5 dto::{
6 abi::v1::CanisterInitPayload,
7 env::EnvBootstrapArgs,
8 subnet::SubnetIdentity,
9 topology::{AppDirectoryArgs, SubnetDirectoryArgs},
10 },
11 ids::CanisterRole,
12};
13use derive_more::{Deref, DerefMut};
14use pocket_ic::{PocketIc, PocketIcBuilder};
15use serde::de::DeserializeOwned;
16
17const INSTALL_CYCLES: u128 = 500 * TC;
18
19#[must_use]
28pub fn pic() -> Pic {
29 PicBuilder::new().with_application_subnet().build()
30}
31
32pub struct PicBuilder(PocketIcBuilder);
42
43#[allow(clippy::new_without_default)]
44impl PicBuilder {
45 #[must_use]
47 pub fn new() -> Self {
48 Self(PocketIcBuilder::new())
49 }
50
51 #[must_use]
53 pub fn with_application_subnet(mut self) -> Self {
54 self.0 = self.0.with_application_subnet();
55 self
56 }
57
58 #[must_use]
60 pub fn with_nns_subnet(mut self) -> Self {
61 self.0 = self.0.with_nns_subnet();
62 self
63 }
64
65 #[must_use]
67 pub fn build(self) -> Pic {
68 Pic(self.0.build())
69 }
70}
71
72#[derive(Deref, DerefMut)]
80pub struct Pic(PocketIc);
81
82impl Pic {
83 pub fn create_and_install_root_canister(&self, wasm: Vec<u8>) -> Result<Principal, Error> {
85 let init_bytes = install_root_args()?;
86
87 Ok(self.create_funded_and_install(wasm, init_bytes))
88 }
89
90 pub fn create_and_install_canister(
94 &self,
95 role: CanisterRole,
96 wasm: Vec<u8>,
97 ) -> Result<Principal, Error> {
98 let init_bytes = install_args(role)?;
99
100 Ok(self.create_funded_and_install(wasm, init_bytes))
101 }
102
103 fn create_funded_and_install(&self, wasm: Vec<u8>, init_bytes: Vec<u8>) -> Principal {
104 let canister_id = self.create_canister();
105 self.add_cycles(canister_id, INSTALL_CYCLES);
106 self.0.install_canister(canister_id, wasm, init_bytes, None);
107
108 canister_id
109 }
110
111 pub fn update_call<T, A>(
113 &self,
114 canister_id: Principal,
115 method: &str,
116 args: A,
117 ) -> Result<T, Error>
118 where
119 T: CandidType + DeserializeOwned,
120 A: ArgumentEncoder,
121 {
122 let bytes: Vec<u8> = encode_args(args)
123 .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
124 let result = self
125 .0
126 .update_call(canister_id, Principal::anonymous(), method, bytes)
127 .map_err(|err| {
128 Error::internal(format!(
129 "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
130 ))
131 })?;
132
133 decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
134 }
135
136 pub fn update_call_as<T, A>(
138 &self,
139 canister_id: Principal,
140 caller: Principal,
141 method: &str,
142 args: A,
143 ) -> Result<T, Error>
144 where
145 T: CandidType + DeserializeOwned,
146 A: ArgumentEncoder,
147 {
148 let bytes: Vec<u8> = encode_args(args)
149 .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
150 let result = self
151 .0
152 .update_call(canister_id, caller, method, bytes)
153 .map_err(|err| {
154 Error::internal(format!(
155 "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
156 ))
157 })?;
158
159 decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
160 }
161
162 pub fn query_call<T, A>(
164 &self,
165 canister_id: Principal,
166 method: &str,
167 args: A,
168 ) -> Result<T, Error>
169 where
170 T: CandidType + DeserializeOwned,
171 A: ArgumentEncoder,
172 {
173 let bytes: Vec<u8> = encode_args(args)
174 .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
175 let result = self
176 .0
177 .query_call(canister_id, Principal::anonymous(), method, bytes)
178 .map_err(|err| {
179 Error::internal(format!(
180 "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
181 ))
182 })?;
183
184 decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
185 }
186
187 pub fn query_call_as<T, A>(
189 &self,
190 canister_id: Principal,
191 caller: Principal,
192 method: &str,
193 args: A,
194 ) -> Result<T, Error>
195 where
196 T: CandidType + DeserializeOwned,
197 A: ArgumentEncoder,
198 {
199 let bytes: Vec<u8> = encode_args(args)
200 .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
201 let result = self
202 .0
203 .query_call(canister_id, caller, method, bytes)
204 .map_err(|err| {
205 Error::internal(format!(
206 "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
207 ))
208 })?;
209
210 decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
211 }
212
213 pub fn tick_n(&self, times: usize) {
214 for _ in 0..times {
215 self.tick();
216 }
217 }
218}
219
220fn install_args(role: CanisterRole) -> Result<Vec<u8>, Error> {
234 if role.is_root() {
235 install_root_args()
236 } else {
237 let env = EnvBootstrapArgs {
240 prime_root_pid: None,
241 subnet_role: None,
242 subnet_pid: None,
243 root_pid: None,
244 canister_role: Some(role),
245 parent_pid: None,
246 };
247
248 let payload = CanisterInitPayload {
251 env,
252 app_directory: AppDirectoryArgs(Vec::new()),
253 subnet_directory: SubnetDirectoryArgs(Vec::new()),
254 };
255
256 encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
257 .map_err(|err| Error::internal(format!("encode_args failed: {err}")))
258 }
259}
260
261fn install_root_args() -> Result<Vec<u8>, Error> {
262 encode_one(SubnetIdentity::Manual)
263 .map_err(|err| Error::internal(format!("encode_one failed: {err}")))
264}