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, SubnetRole},
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 let args = if role.is_root() {
265 let subnet_pid = Principal::from_slice(&[0xAA; 29]);
267 encode_one(SubnetIdentity::Manual(subnet_pid))
268 .map_err(|err| PublicError::internal(format!("encode_one failed: {err}")))
269 } else {
270 let root_pid = Principal::from_slice(&[0xBB; 29]);
272 let subnet_pid = Principal::from_slice(&[0xAA; 29]);
273 let env = EnvView {
274 prime_root_pid: Some(root_pid),
275 subnet_role: Some(SubnetRole::PRIME),
276 subnet_pid: Some(subnet_pid),
277 root_pid: Some(root_pid),
278 canister_role: Some(role),
279 parent_pid: Some(root_pid),
280 };
281
282 let payload = CanisterInitPayload {
285 env,
286 app_directory: AppDirectoryView(Vec::new()),
287 subnet_directory: SubnetDirectoryView(Vec::new()),
288 };
289 encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
290 .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))
291 }?;
292
293 Ok(args)
294}
295
296fn install_args_with_directories(
297 role: CanisterRole,
298 app_directory: AppDirectoryView,
299 subnet_directory: SubnetDirectoryView,
300) -> Result<Vec<u8>, PublicError> {
301 let args = if role.is_root() {
302 let subnet_pid = Principal::from_slice(&[0xAA; 29]);
303 encode_one(SubnetIdentity::Manual(subnet_pid))
304 .map_err(|err| PublicError::internal(format!("encode_one failed: {err}")))
305 } else {
306 let root_pid = Principal::from_slice(&[0xBB; 29]);
307 let subnet_pid = Principal::from_slice(&[0xAA; 29]);
308 let env = EnvView {
309 prime_root_pid: Some(root_pid),
310 subnet_role: Some(SubnetRole::PRIME),
311 subnet_pid: Some(subnet_pid),
312 root_pid: Some(root_pid),
313 canister_role: Some(role),
314 parent_pid: Some(root_pid),
315 };
316 let payload = CanisterInitPayload {
317 env,
318 app_directory,
319 subnet_directory,
320 };
321 encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
322 .map_err(|err| PublicError::internal(format!("encode_args failed: {err}")))
323 }?;
324
325 Ok(args)
326}