1use candid::{CandidType, Principal, decode_one, encode_args, encode_one, utils::ArgumentEncoder};
2use canic::{
3 PublicError,
4 cdk::types::TC,
5 core::{
6 dto::{
7 abi::v1::CanisterInitPayload,
8 env::EnvView,
9 subnet::SubnetIdentity,
10 topology::{AppDirectoryView, SubnetDirectoryView},
11 },
12 ids::CanisterRole,
13 },
14};
15use derive_more::{Deref, DerefMut};
16use pocket_ic::{PocketIc, PocketIcBuilder};
17use serde::de::DeserializeOwned;
18
19const INSTALL_CYCLES: u128 = 500 * TC;
20
21#[must_use]
30pub fn pic() -> Pic {
31 PicBuilder::new().with_application_subnet().build()
32}
33
34pub struct PicBuilder(PocketIcBuilder);
44
45#[allow(clippy::new_without_default)]
46impl PicBuilder {
47 #[must_use]
49 pub fn new() -> Self {
50 Self(PocketIcBuilder::new())
51 }
52
53 #[must_use]
55 pub fn with_application_subnet(mut self) -> Self {
56 self.0 = self.0.with_application_subnet();
57 self
58 }
59
60 #[must_use]
62 pub fn with_nns_subnet(mut self) -> Self {
63 self.0 = self.0.with_nns_subnet();
64 self
65 }
66
67 #[must_use]
69 pub fn build(self) -> Pic {
70 Pic(self.0.build())
71 }
72}
73
74#[derive(Deref, DerefMut)]
82pub struct Pic(PocketIc);
83
84impl Pic {
85 pub fn create_and_install_root_canister(
87 &self,
88 wasm: Vec<u8>,
89 ) -> Result<Principal, PublicError> {
90 let init_bytes = install_root_args()?;
91
92 Ok(self.create_funded_and_install(wasm, init_bytes))
93 }
94
95 pub fn create_and_install_canister(
99 &self,
100 role: CanisterRole,
101 wasm: Vec<u8>,
102 ) -> Result<Principal, PublicError> {
103 let init_bytes = install_args(role)?;
104
105 Ok(self.create_funded_and_install(wasm, init_bytes))
106 }
107
108 pub fn create_and_install_canister_with_directories(
113 &self,
114 role: CanisterRole,
115 wasm: Vec<u8>,
116 app_directory: AppDirectoryView,
117 subnet_directory: SubnetDirectoryView,
118 ) -> Result<Principal, PublicError> {
119 let init_bytes = install_args_with_directories(role, app_directory, subnet_directory)?;
120
121 Ok(self.create_funded_and_install(wasm, init_bytes))
122 }
123
124 fn create_funded_and_install(&self, wasm: Vec<u8>, init_bytes: Vec<u8>) -> Principal {
125 let canister_id = self.create_canister();
126 self.add_cycles(canister_id, INSTALL_CYCLES);
127 self.0.install_canister(canister_id, wasm, init_bytes, None);
128
129 canister_id
130 }
131
132 pub fn update_call<T, A>(
134 &self,
135 canister_id: Principal,
136 method: &str,
137 args: A,
138 ) -> Result<T, PublicError>
139 where
140 T: CandidType + DeserializeOwned,
141 A: ArgumentEncoder,
142 {
143 let bytes: Vec<u8> = encode_args(args)
144 .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))?;
145 let result = self
146 .0
147 .update_call(canister_id, Principal::anonymous(), method, bytes)
148 .map_err(|err| {
149 PublicError::internal(format!(
150 "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
151 ))
152 })?;
153
154 decode_one(&result)
155 .map_err(|err| PublicError::internal(format!("decode_one failed: {err}")))
156 }
157
158 pub fn update_call_as<T, A>(
160 &self,
161 canister_id: Principal,
162 caller: Principal,
163 method: &str,
164 args: A,
165 ) -> Result<T, PublicError>
166 where
167 T: CandidType + DeserializeOwned,
168 A: ArgumentEncoder,
169 {
170 let bytes: Vec<u8> = encode_args(args)
171 .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))?;
172 let result = self
173 .0
174 .update_call(canister_id, caller, method, bytes)
175 .map_err(|err| {
176 PublicError::internal(format!(
177 "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
178 ))
179 })?;
180
181 decode_one(&result)
182 .map_err(|err| PublicError::internal(format!("decode_one failed: {err}")))
183 }
184
185 pub fn query_call<T, A>(
187 &self,
188 canister_id: Principal,
189 method: &str,
190 args: A,
191 ) -> Result<T, PublicError>
192 where
193 T: CandidType + DeserializeOwned,
194 A: ArgumentEncoder,
195 {
196 let bytes: Vec<u8> = encode_args(args)
197 .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))?;
198 let result = self
199 .0
200 .query_call(canister_id, Principal::anonymous(), method, bytes)
201 .map_err(|err| {
202 PublicError::internal(format!(
203 "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
204 ))
205 })?;
206
207 decode_one(&result)
208 .map_err(|err| PublicError::internal(format!("decode_one failed: {err}")))
209 }
210
211 pub fn query_call_as<T, A>(
213 &self,
214 canister_id: Principal,
215 caller: Principal,
216 method: &str,
217 args: A,
218 ) -> Result<T, PublicError>
219 where
220 T: CandidType + DeserializeOwned,
221 A: ArgumentEncoder,
222 {
223 let bytes: Vec<u8> = encode_args(args)
224 .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))?;
225 let result = self
226 .0
227 .query_call(canister_id, caller, method, bytes)
228 .map_err(|err| {
229 PublicError::internal(format!(
230 "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
231 ))
232 })?;
233
234 decode_one(&result)
235 .map_err(|err| PublicError::internal(format!("decode_one failed: {err}")))
236 }
237
238 pub fn tick_n(&self, times: usize) {
239 for _ in 0..times {
240 self.tick();
241 }
242 }
243}
244
245fn install_args(role: CanisterRole) -> Result<Vec<u8>, PublicError> {
259 if role.is_root() {
260 install_root_args()
261 } else {
262 let env = EnvView {
265 prime_root_pid: None,
266 subnet_role: None,
267 subnet_pid: None,
268 root_pid: None,
269 canister_role: Some(role),
270 parent_pid: None,
271 };
272
273 let payload = CanisterInitPayload {
276 env,
277 app_directory: AppDirectoryView(Vec::new()),
278 subnet_directory: SubnetDirectoryView(Vec::new()),
279 };
280
281 encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
282 .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))
283 }
284}
285
286fn install_root_args() -> Result<Vec<u8>, PublicError> {
287 encode_one(SubnetIdentity::Manual)
288 .map_err(|err| PublicError::internal(format!("encode_one failed: {err}")))
289}
290
291fn install_args_with_directories(
292 role: CanisterRole,
293 app_directory: AppDirectoryView,
294 subnet_directory: SubnetDirectoryView,
295) -> Result<Vec<u8>, PublicError> {
296 if role.is_root() {
297 encode_one(SubnetIdentity::Manual)
300 .map_err(|err| PublicError::internal(format!("encode_one failed: {err}")))
301 } else {
302 let env = EnvView {
304 prime_root_pid: None,
305 subnet_role: None,
306 subnet_pid: None,
307 root_pid: None,
308 canister_role: Some(role),
309 parent_pid: None,
310 };
311
312 let payload = CanisterInitPayload {
313 env,
314 app_directory,
315 subnet_directory,
316 };
317
318 encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
319 .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))
320 }
321}