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 pocket_ic::{PocketIc, PocketIcBuilder};
14use serde::de::DeserializeOwned;
15use std::{
16 ops::{Deref, DerefMut},
17 panic::{AssertUnwindSafe, catch_unwind},
18 sync::{Mutex, MutexGuard},
19};
20
21const INSTALL_CYCLES: u128 = 500 * TC;
22static PIC_BUILD_SERIAL: Mutex<()> = Mutex::new(());
23
24#[must_use]
35pub fn pic() -> Pic {
36 PicBuilder::new().with_application_subnet().build()
37}
38
39pub struct PicBuilder(PocketIcBuilder);
49
50#[expect(clippy::new_without_default)]
51impl PicBuilder {
52 #[must_use]
54 pub fn new() -> Self {
55 Self(PocketIcBuilder::new())
56 }
57
58 #[must_use]
60 pub fn with_application_subnet(mut self) -> Self {
61 self.0 = self.0.with_application_subnet();
62 self
63 }
64
65 #[must_use]
67 pub fn with_nns_subnet(mut self) -> Self {
68 self.0 = self.0.with_nns_subnet();
69 self
70 }
71
72 #[must_use]
74 pub fn build(self) -> Pic {
75 let serial_guard = PIC_BUILD_SERIAL
79 .lock()
80 .unwrap_or_else(std::sync::PoisonError::into_inner);
81
82 Pic {
83 inner: self.0.build(),
84 _serial_guard: serial_guard,
85 }
86 }
87}
88
89pub struct Pic {
100 inner: PocketIc,
101 _serial_guard: MutexGuard<'static, ()>,
102}
103
104impl Pic {
105 pub fn create_and_install_root_canister(&self, wasm: Vec<u8>) -> Result<Principal, Error> {
107 let init_bytes = install_root_args()?;
108
109 Ok(self.create_funded_and_install(wasm, init_bytes))
110 }
111
112 pub fn create_and_install_canister(
116 &self,
117 role: CanisterRole,
118 wasm: Vec<u8>,
119 ) -> Result<Principal, Error> {
120 let init_bytes = install_args(role)?;
121
122 Ok(self.create_funded_and_install(wasm, init_bytes))
123 }
124
125 fn create_funded_and_install(&self, wasm: Vec<u8>, init_bytes: Vec<u8>) -> Principal {
126 let canister_id = self.create_canister();
127 self.add_cycles(canister_id, INSTALL_CYCLES);
128
129 let install = catch_unwind(AssertUnwindSafe(|| {
130 self.inner
131 .install_canister(canister_id, wasm, init_bytes, None);
132 }));
133 if let Err(err) = install {
134 eprintln!("install_canister trapped for {canister_id}");
135 if let Ok(status) = self.inner.canister_status(canister_id, None) {
136 eprintln!("canister_status for {canister_id}: {status:?}");
137 }
138 if let Ok(logs) = self
139 .inner
140 .fetch_canister_logs(canister_id, Principal::anonymous())
141 {
142 for record in logs {
143 eprintln!("canister_log {canister_id}: {record:?}");
144 }
145 }
146 std::panic::resume_unwind(err);
147 }
148
149 canister_id
150 }
151
152 pub fn update_call<T, A>(
154 &self,
155 canister_id: Principal,
156 method: &str,
157 args: A,
158 ) -> Result<T, Error>
159 where
160 T: CandidType + DeserializeOwned,
161 A: ArgumentEncoder,
162 {
163 let bytes: Vec<u8> = encode_args(args)
164 .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
165 let result = self
166 .inner
167 .update_call(canister_id, Principal::anonymous(), method, bytes)
168 .map_err(|err| {
169 Error::internal(format!(
170 "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
171 ))
172 })?;
173
174 decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
175 }
176
177 pub fn update_call_as<T, A>(
179 &self,
180 canister_id: Principal,
181 caller: Principal,
182 method: &str,
183 args: A,
184 ) -> Result<T, Error>
185 where
186 T: CandidType + DeserializeOwned,
187 A: ArgumentEncoder,
188 {
189 let bytes: Vec<u8> = encode_args(args)
190 .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
191 let result = self
192 .inner
193 .update_call(canister_id, caller, method, bytes)
194 .map_err(|err| {
195 Error::internal(format!(
196 "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
197 ))
198 })?;
199
200 decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
201 }
202
203 pub fn query_call<T, A>(
205 &self,
206 canister_id: Principal,
207 method: &str,
208 args: A,
209 ) -> Result<T, Error>
210 where
211 T: CandidType + DeserializeOwned,
212 A: ArgumentEncoder,
213 {
214 let bytes: Vec<u8> = encode_args(args)
215 .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
216 let result = self
217 .inner
218 .query_call(canister_id, Principal::anonymous(), method, bytes)
219 .map_err(|err| {
220 Error::internal(format!(
221 "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
222 ))
223 })?;
224
225 decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
226 }
227
228 pub fn query_call_as<T, A>(
230 &self,
231 canister_id: Principal,
232 caller: Principal,
233 method: &str,
234 args: A,
235 ) -> Result<T, Error>
236 where
237 T: CandidType + DeserializeOwned,
238 A: ArgumentEncoder,
239 {
240 let bytes: Vec<u8> = encode_args(args)
241 .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
242 let result = self
243 .inner
244 .query_call(canister_id, caller, method, bytes)
245 .map_err(|err| {
246 Error::internal(format!(
247 "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
248 ))
249 })?;
250
251 decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
252 }
253
254 pub fn tick_n(&self, times: usize) {
255 for _ in 0..times {
256 self.tick();
257 }
258 }
259}
260
261impl Deref for Pic {
262 type Target = PocketIc;
263
264 fn deref(&self) -> &Self::Target {
265 &self.inner
266 }
267}
268
269impl DerefMut for Pic {
270 fn deref_mut(&mut self) -> &mut Self::Target {
271 &mut self.inner
272 }
273}
274
275fn install_args(role: CanisterRole) -> Result<Vec<u8>, Error> {
289 if role.is_root() {
290 install_root_args()
291 } else {
292 let env = EnvBootstrapArgs {
295 prime_root_pid: None,
296 subnet_role: None,
297 subnet_pid: None,
298 root_pid: None,
299 canister_role: Some(role),
300 parent_pid: None,
301 };
302
303 let payload = CanisterInitPayload {
306 env,
307 app_directory: AppDirectoryArgs(Vec::new()),
308 subnet_directory: SubnetDirectoryArgs(Vec::new()),
309 };
310
311 encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
312 .map_err(|err| Error::internal(format!("encode_args failed: {err}")))
313 }
314}
315
316fn install_root_args() -> Result<Vec<u8>, Error> {
317 encode_one(SubnetIdentity::Manual)
318 .map_err(|err| Error::internal(format!("encode_one failed: {err}")))
319}