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;
17use std::sync::OnceLock;
18
19static PIC: OnceLock<Pic> = OnceLock::new();
36
37#[must_use]
43pub fn pic() -> &'static Pic {
44 PIC.get_or_init(|| PicBuilder::new().with_application_subnet().build())
45}
46
47pub struct PicBuilder(PocketIcBuilder);
57
58#[allow(clippy::new_without_default)]
59impl PicBuilder {
60 #[must_use]
62 pub fn new() -> Self {
63 Self(PocketIcBuilder::new())
64 }
65
66 #[must_use]
68 pub fn with_application_subnet(mut self) -> Self {
69 self.0 = self.0.with_application_subnet();
70 self
71 }
72
73 #[must_use]
75 pub fn with_nns_subnet(mut self) -> Self {
76 self.0 = self.0.with_nns_subnet();
77 self
78 }
79
80 #[must_use]
82 pub fn build(self) -> Pic {
83 Pic(self.0.build())
84 }
85}
86
87#[derive(Deref, DerefMut)]
95pub struct Pic(PocketIc);
96
97impl Pic {
98 pub fn create_and_install_canister(
102 &self,
103 role: CanisterRole,
104 wasm: Vec<u8>,
105 ) -> Result<Principal, PublicError> {
106 let canister_id = self.create_canister();
108 self.add_cycles(canister_id, 1_000_000_000_000);
109
110 let init_bytes = install_args(role)?;
112 self.0.install_canister(canister_id, wasm, init_bytes, None);
113
114 Ok(canister_id)
115 }
116
117 pub fn create_and_install_canister_with_directories(
122 &self,
123 role: CanisterRole,
124 wasm: Vec<u8>,
125 app_directory: AppDirectoryView,
126 subnet_directory: SubnetDirectoryView,
127 ) -> Result<Principal, PublicError> {
128 let canister_id = self.create_canister();
129 self.add_cycles(canister_id, 1_000_000_000_000);
130
131 let init_bytes = install_args_with_directories(role, app_directory, subnet_directory)?;
132 self.0.install_canister(canister_id, wasm, init_bytes, None);
133
134 Ok(canister_id)
135 }
136
137 pub fn update_call<T, A>(
139 &self,
140 canister_id: Principal,
141 method: &str,
142 args: A,
143 ) -> Result<T, PublicError>
144 where
145 T: CandidType + DeserializeOwned,
146 A: ArgumentEncoder,
147 {
148 let bytes: Vec<u8> = encode_args(args)
149 .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))?;
150 let result = self
151 .0
152 .update_call(canister_id, Principal::anonymous(), method, bytes)
153 .map_err(|err| {
154 PublicError::internal(format!(
155 "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
156 ))
157 })?;
158
159 decode_one(&result)
160 .map_err(|err| PublicError::internal(format!("decode_one failed: {err}")))
161 }
162
163 pub fn update_call_as<T, A>(
165 &self,
166 canister_id: Principal,
167 caller: Principal,
168 method: &str,
169 args: A,
170 ) -> Result<T, PublicError>
171 where
172 T: CandidType + DeserializeOwned,
173 A: ArgumentEncoder,
174 {
175 let bytes: Vec<u8> = encode_args(args)
176 .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))?;
177 let result = self
178 .0
179 .update_call(canister_id, caller, method, bytes)
180 .map_err(|err| {
181 PublicError::internal(format!(
182 "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
183 ))
184 })?;
185
186 decode_one(&result)
187 .map_err(|err| PublicError::internal(format!("decode_one failed: {err}")))
188 }
189
190 pub fn query_call<T, A>(
192 &self,
193 canister_id: Principal,
194 method: &str,
195 args: A,
196 ) -> Result<T, PublicError>
197 where
198 T: CandidType + DeserializeOwned,
199 A: ArgumentEncoder,
200 {
201 let bytes: Vec<u8> = encode_args(args)
202 .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))?;
203 let result = self
204 .0
205 .query_call(canister_id, Principal::anonymous(), method, bytes)
206 .map_err(|err| {
207 PublicError::internal(format!(
208 "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
209 ))
210 })?;
211
212 decode_one(&result)
213 .map_err(|err| PublicError::internal(format!("decode_one failed: {err}")))
214 }
215
216 pub fn query_call_as<T, A>(
218 &self,
219 canister_id: Principal,
220 caller: Principal,
221 method: &str,
222 args: A,
223 ) -> Result<T, PublicError>
224 where
225 T: CandidType + DeserializeOwned,
226 A: ArgumentEncoder,
227 {
228 let bytes: Vec<u8> = encode_args(args)
229 .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))?;
230 let result = self
231 .0
232 .query_call(canister_id, caller, method, bytes)
233 .map_err(|err| {
234 PublicError::internal(format!(
235 "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
236 ))
237 })?;
238
239 decode_one(&result)
240 .map_err(|err| PublicError::internal(format!("decode_one failed: {err}")))
241 }
242
243 pub fn tick_n(&self, times: usize) {
244 for _ in 0..times {
245 self.tick();
246 }
247 }
248}
249
250fn install_args(role: CanisterRole) -> Result<Vec<u8>, PublicError> {
264 if role.is_root() {
265 encode_one(SubnetIdentity::Manual)
268 .map_err(|err| PublicError::internal(format!("encode_one failed: {err}")))
269 } else {
270 let env = EnvView {
273 prime_root_pid: None,
274 subnet_role: None,
275 subnet_pid: None,
276 root_pid: None,
277 canister_role: Some(role),
278 parent_pid: None,
279 };
280
281 let payload = CanisterInitPayload {
284 env,
285 app_directory: AppDirectoryView(Vec::new()),
286 subnet_directory: SubnetDirectoryView(Vec::new()),
287 };
288
289 encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
290 .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))
291 }
292}
293
294fn install_args_with_directories(
295 role: CanisterRole,
296 app_directory: AppDirectoryView,
297 subnet_directory: SubnetDirectoryView,
298) -> Result<Vec<u8>, PublicError> {
299 if role.is_root() {
300 encode_one(SubnetIdentity::Manual)
303 .map_err(|err| PublicError::internal(format!("encode_one failed: {err}")))
304 } else {
305 let env = EnvView {
307 prime_root_pid: None,
308 subnet_role: None,
309 subnet_pid: None,
310 root_pid: None,
311 canister_role: Some(role),
312 parent_pid: None,
313 };
314
315 let payload = CanisterInitPayload {
316 env,
317 app_directory,
318 subnet_directory,
319 };
320
321 encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
322 .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))
323 }
324}