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 sync::{Mutex, MutexGuard},
18};
19
20const INSTALL_CYCLES: u128 = 500 * TC;
21static PIC_BUILD_SERIAL: Mutex<()> = Mutex::new(());
22
23#[must_use]
34pub fn pic() -> Pic {
35 PicBuilder::new().with_application_subnet().build()
36}
37
38pub struct PicBuilder(PocketIcBuilder);
48
49#[expect(clippy::new_without_default)]
50impl PicBuilder {
51 #[must_use]
53 pub fn new() -> Self {
54 Self(PocketIcBuilder::new())
55 }
56
57 #[must_use]
59 pub fn with_application_subnet(mut self) -> Self {
60 self.0 = self.0.with_application_subnet();
61 self
62 }
63
64 #[must_use]
66 pub fn with_nns_subnet(mut self) -> Self {
67 self.0 = self.0.with_nns_subnet();
68 self
69 }
70
71 #[must_use]
73 pub fn build(self) -> Pic {
74 let serial_guard = PIC_BUILD_SERIAL
78 .lock()
79 .unwrap_or_else(std::sync::PoisonError::into_inner);
80
81 Pic {
82 inner: self.0.build(),
83 _serial_guard: serial_guard,
84 }
85 }
86}
87
88pub struct Pic {
99 inner: PocketIc,
100 _serial_guard: MutexGuard<'static, ()>,
101}
102
103impl Deref for Pic {
104 type Target = PocketIc;
105
106 fn deref(&self) -> &Self::Target {
107 &self.inner
108 }
109}
110
111impl DerefMut for Pic {
112 fn deref_mut(&mut self) -> &mut Self::Target {
113 &mut self.inner
114 }
115}
116
117impl Pic {
118 pub fn create_and_install_root_canister(&self, wasm: Vec<u8>) -> Result<Principal, Error> {
120 let init_bytes = install_root_args()?;
121
122 Ok(self.create_funded_and_install(wasm, init_bytes))
123 }
124
125 pub fn create_and_install_canister(
129 &self,
130 role: CanisterRole,
131 wasm: Vec<u8>,
132 ) -> Result<Principal, Error> {
133 let init_bytes = install_args(role)?;
134
135 Ok(self.create_funded_and_install(wasm, init_bytes))
136 }
137
138 fn create_funded_and_install(&self, wasm: Vec<u8>, init_bytes: Vec<u8>) -> Principal {
139 let canister_id = self.create_canister();
140 self.add_cycles(canister_id, INSTALL_CYCLES);
141 self.inner
142 .install_canister(canister_id, wasm, init_bytes, None);
143
144 canister_id
145 }
146
147 pub fn update_call<T, A>(
149 &self,
150 canister_id: Principal,
151 method: &str,
152 args: A,
153 ) -> Result<T, Error>
154 where
155 T: CandidType + DeserializeOwned,
156 A: ArgumentEncoder,
157 {
158 let bytes: Vec<u8> = encode_args(args)
159 .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
160 let result = self
161 .inner
162 .update_call(canister_id, Principal::anonymous(), method, bytes)
163 .map_err(|err| {
164 Error::internal(format!(
165 "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
166 ))
167 })?;
168
169 decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
170 }
171
172 pub fn update_call_as<T, A>(
174 &self,
175 canister_id: Principal,
176 caller: Principal,
177 method: &str,
178 args: A,
179 ) -> Result<T, Error>
180 where
181 T: CandidType + DeserializeOwned,
182 A: ArgumentEncoder,
183 {
184 let bytes: Vec<u8> = encode_args(args)
185 .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
186 let result = self
187 .inner
188 .update_call(canister_id, caller, method, bytes)
189 .map_err(|err| {
190 Error::internal(format!(
191 "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
192 ))
193 })?;
194
195 decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
196 }
197
198 pub fn query_call<T, A>(
200 &self,
201 canister_id: Principal,
202 method: &str,
203 args: A,
204 ) -> Result<T, Error>
205 where
206 T: CandidType + DeserializeOwned,
207 A: ArgumentEncoder,
208 {
209 let bytes: Vec<u8> = encode_args(args)
210 .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
211 let result = self
212 .inner
213 .query_call(canister_id, Principal::anonymous(), method, bytes)
214 .map_err(|err| {
215 Error::internal(format!(
216 "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
217 ))
218 })?;
219
220 decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
221 }
222
223 pub fn query_call_as<T, A>(
225 &self,
226 canister_id: Principal,
227 caller: Principal,
228 method: &str,
229 args: A,
230 ) -> Result<T, Error>
231 where
232 T: CandidType + DeserializeOwned,
233 A: ArgumentEncoder,
234 {
235 let bytes: Vec<u8> = encode_args(args)
236 .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
237 let result = self
238 .inner
239 .query_call(canister_id, caller, method, bytes)
240 .map_err(|err| {
241 Error::internal(format!(
242 "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
243 ))
244 })?;
245
246 decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
247 }
248
249 pub fn tick_n(&self, times: usize) {
250 for _ in 0..times {
251 self.tick();
252 }
253 }
254}
255
256fn install_args(role: CanisterRole) -> Result<Vec<u8>, Error> {
270 if role.is_root() {
271 install_root_args()
272 } else {
273 let env = EnvBootstrapArgs {
276 prime_root_pid: None,
277 subnet_role: None,
278 subnet_pid: None,
279 root_pid: None,
280 canister_role: Some(role),
281 parent_pid: None,
282 };
283
284 let payload = CanisterInitPayload {
287 env,
288 app_directory: AppDirectoryArgs(Vec::new()),
289 subnet_directory: SubnetDirectoryArgs(Vec::new()),
290 };
291
292 encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
293 .map_err(|err| Error::internal(format!("encode_args failed: {err}")))
294 }
295}
296
297fn install_root_args() -> Result<Vec<u8>, Error> {
298 encode_one(SubnetIdentity::Manual)
299 .map_err(|err| Error::internal(format!("encode_one failed: {err}")))
300}