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