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 Pic {
104 pub fn create_and_install_root_canister(&self, wasm: Vec<u8>) -> Result<Principal, Error> {
106 let init_bytes = install_root_args()?;
107
108 Ok(self.create_funded_and_install(wasm, init_bytes))
109 }
110
111 pub fn create_and_install_canister(
115 &self,
116 role: CanisterRole,
117 wasm: Vec<u8>,
118 ) -> Result<Principal, Error> {
119 let init_bytes = install_args(role)?;
120
121 Ok(self.create_funded_and_install(wasm, init_bytes))
122 }
123
124 fn create_funded_and_install(&self, wasm: Vec<u8>, init_bytes: Vec<u8>) -> Principal {
125 let canister_id = self.create_canister();
126 self.add_cycles(canister_id, INSTALL_CYCLES);
127 self.inner
128 .install_canister(canister_id, wasm, init_bytes, None);
129
130 canister_id
131 }
132
133 pub fn update_call<T, A>(
135 &self,
136 canister_id: Principal,
137 method: &str,
138 args: A,
139 ) -> Result<T, Error>
140 where
141 T: CandidType + DeserializeOwned,
142 A: ArgumentEncoder,
143 {
144 let bytes: Vec<u8> = encode_args(args)
145 .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
146 let result = self
147 .inner
148 .update_call(canister_id, Principal::anonymous(), method, bytes)
149 .map_err(|err| {
150 Error::internal(format!(
151 "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
152 ))
153 })?;
154
155 decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
156 }
157
158 pub fn update_call_as<T, A>(
160 &self,
161 canister_id: Principal,
162 caller: Principal,
163 method: &str,
164 args: A,
165 ) -> Result<T, Error>
166 where
167 T: CandidType + DeserializeOwned,
168 A: ArgumentEncoder,
169 {
170 let bytes: Vec<u8> = encode_args(args)
171 .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
172 let result = self
173 .inner
174 .update_call(canister_id, caller, method, bytes)
175 .map_err(|err| {
176 Error::internal(format!(
177 "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
178 ))
179 })?;
180
181 decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
182 }
183
184 pub fn query_call<T, A>(
186 &self,
187 canister_id: Principal,
188 method: &str,
189 args: A,
190 ) -> Result<T, Error>
191 where
192 T: CandidType + DeserializeOwned,
193 A: ArgumentEncoder,
194 {
195 let bytes: Vec<u8> = encode_args(args)
196 .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
197 let result = self
198 .inner
199 .query_call(canister_id, Principal::anonymous(), method, bytes)
200 .map_err(|err| {
201 Error::internal(format!(
202 "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
203 ))
204 })?;
205
206 decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
207 }
208
209 pub fn query_call_as<T, A>(
211 &self,
212 canister_id: Principal,
213 caller: Principal,
214 method: &str,
215 args: A,
216 ) -> Result<T, Error>
217 where
218 T: CandidType + DeserializeOwned,
219 A: ArgumentEncoder,
220 {
221 let bytes: Vec<u8> = encode_args(args)
222 .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
223 let result = self
224 .inner
225 .query_call(canister_id, caller, method, bytes)
226 .map_err(|err| {
227 Error::internal(format!(
228 "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
229 ))
230 })?;
231
232 decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
233 }
234
235 pub fn tick_n(&self, times: usize) {
236 for _ in 0..times {
237 self.tick();
238 }
239 }
240}
241
242impl Deref for Pic {
243 type Target = PocketIc;
244
245 fn deref(&self) -> &Self::Target {
246 &self.inner
247 }
248}
249
250impl DerefMut for Pic {
251 fn deref_mut(&mut self) -> &mut Self::Target {
252 &mut self.inner
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}