1use candid::{CandidType, Principal, decode_one, encode_args, encode_one, utils::ArgumentEncoder};
2use canic::{
3 PublicError,
4 core::{
5 dto::{
6 abi::v1::CanisterInitPayload,
7 env::EnvView,
8 subnet::SubnetIdentity,
9 topology::{AppDirectoryView, SubnetDirectoryView},
10 },
11 ids::CanisterRole,
12 },
13};
14use derive_more::{Deref, DerefMut};
15use pocket_ic::{PocketIc, PocketIcBuilder};
16use serde::de::DeserializeOwned;
17
18#[must_use]
27pub fn pic() -> Pic {
28 PicBuilder::new().with_application_subnet().build()
29}
30
31pub struct PicBuilder(PocketIcBuilder);
41
42#[allow(clippy::new_without_default)]
43impl PicBuilder {
44 #[must_use]
46 pub fn new() -> Self {
47 Self(PocketIcBuilder::new())
48 }
49
50 #[must_use]
52 pub fn with_application_subnet(mut self) -> Self {
53 self.0 = self.0.with_application_subnet();
54 self
55 }
56
57 #[must_use]
59 pub fn with_nns_subnet(mut self) -> Self {
60 self.0 = self.0.with_nns_subnet();
61 self
62 }
63
64 #[must_use]
66 pub fn build(self) -> Pic {
67 Pic(self.0.build())
68 }
69}
70
71#[derive(Deref, DerefMut)]
79pub struct Pic(PocketIc);
80
81impl Pic {
82 pub fn create_and_install_canister(
86 &self,
87 role: CanisterRole,
88 wasm: Vec<u8>,
89 ) -> Result<Principal, PublicError> {
90 let canister_id = self.create_canister();
92 self.add_cycles(canister_id, 1_000_000_000_000);
93
94 let init_bytes = install_args(role)?;
96 self.0.install_canister(canister_id, wasm, init_bytes, None);
97
98 Ok(canister_id)
99 }
100
101 pub fn create_and_install_canister_with_directories(
106 &self,
107 role: CanisterRole,
108 wasm: Vec<u8>,
109 app_directory: AppDirectoryView,
110 subnet_directory: SubnetDirectoryView,
111 ) -> Result<Principal, PublicError> {
112 let canister_id = self.create_canister();
113 self.add_cycles(canister_id, 1_000_000_000_000);
114
115 let init_bytes = install_args_with_directories(role, app_directory, subnet_directory)?;
116 self.0.install_canister(canister_id, wasm, init_bytes, None);
117
118 Ok(canister_id)
119 }
120
121 pub fn update_call<T, A>(
123 &self,
124 canister_id: Principal,
125 method: &str,
126 args: A,
127 ) -> Result<T, PublicError>
128 where
129 T: CandidType + DeserializeOwned,
130 A: ArgumentEncoder,
131 {
132 let bytes: Vec<u8> = encode_args(args)
133 .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))?;
134 let result = self
135 .0
136 .update_call(canister_id, Principal::anonymous(), method, bytes)
137 .map_err(|err| {
138 PublicError::internal(format!(
139 "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
140 ))
141 })?;
142
143 decode_one(&result)
144 .map_err(|err| PublicError::internal(format!("decode_one failed: {err}")))
145 }
146
147 pub fn update_call_as<T, A>(
149 &self,
150 canister_id: Principal,
151 caller: Principal,
152 method: &str,
153 args: A,
154 ) -> Result<T, PublicError>
155 where
156 T: CandidType + DeserializeOwned,
157 A: ArgumentEncoder,
158 {
159 let bytes: Vec<u8> = encode_args(args)
160 .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))?;
161 let result = self
162 .0
163 .update_call(canister_id, caller, method, bytes)
164 .map_err(|err| {
165 PublicError::internal(format!(
166 "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
167 ))
168 })?;
169
170 decode_one(&result)
171 .map_err(|err| PublicError::internal(format!("decode_one failed: {err}")))
172 }
173
174 pub fn query_call<T, A>(
176 &self,
177 canister_id: Principal,
178 method: &str,
179 args: A,
180 ) -> Result<T, PublicError>
181 where
182 T: CandidType + DeserializeOwned,
183 A: ArgumentEncoder,
184 {
185 let bytes: Vec<u8> = encode_args(args)
186 .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))?;
187 let result = self
188 .0
189 .query_call(canister_id, Principal::anonymous(), method, bytes)
190 .map_err(|err| {
191 PublicError::internal(format!(
192 "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
193 ))
194 })?;
195
196 decode_one(&result)
197 .map_err(|err| PublicError::internal(format!("decode_one failed: {err}")))
198 }
199
200 pub fn query_call_as<T, A>(
202 &self,
203 canister_id: Principal,
204 caller: Principal,
205 method: &str,
206 args: A,
207 ) -> Result<T, PublicError>
208 where
209 T: CandidType + DeserializeOwned,
210 A: ArgumentEncoder,
211 {
212 let bytes: Vec<u8> = encode_args(args)
213 .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))?;
214 let result = self
215 .0
216 .query_call(canister_id, caller, method, bytes)
217 .map_err(|err| {
218 PublicError::internal(format!(
219 "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
220 ))
221 })?;
222
223 decode_one(&result)
224 .map_err(|err| PublicError::internal(format!("decode_one failed: {err}")))
225 }
226
227 pub fn tick_n(&self, times: usize) {
228 for _ in 0..times {
229 self.tick();
230 }
231 }
232}
233
234fn install_args(role: CanisterRole) -> Result<Vec<u8>, PublicError> {
248 if role.is_root() {
249 encode_one(SubnetIdentity::Manual)
252 .map_err(|err| PublicError::internal(format!("encode_one failed: {err}")))
253 } else {
254 let env = EnvView {
257 prime_root_pid: None,
258 subnet_role: None,
259 subnet_pid: None,
260 root_pid: None,
261 canister_role: Some(role),
262 parent_pid: None,
263 };
264
265 let payload = CanisterInitPayload {
268 env,
269 app_directory: AppDirectoryView(Vec::new()),
270 subnet_directory: SubnetDirectoryView(Vec::new()),
271 };
272
273 encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
274 .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))
275 }
276}
277
278fn install_args_with_directories(
279 role: CanisterRole,
280 app_directory: AppDirectoryView,
281 subnet_directory: SubnetDirectoryView,
282) -> Result<Vec<u8>, PublicError> {
283 if role.is_root() {
284 encode_one(SubnetIdentity::Manual)
287 .map_err(|err| PublicError::internal(format!("encode_one failed: {err}")))
288 } else {
289 let env = EnvView {
291 prime_root_pid: None,
292 subnet_role: None,
293 subnet_pid: None,
294 root_pid: None,
295 canister_role: Some(role),
296 parent_pid: None,
297 };
298
299 let payload = CanisterInitPayload {
300 env,
301 app_directory,
302 subnet_directory,
303 };
304
305 encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
306 .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))
307 }
308}