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::EnvView,
8 subnet::SubnetIdentity,
9 topology::{AppDirectoryView, SubnetDirectoryView},
10 },
11 ids::CanisterRole,
12};
13use derive_more::{Deref, DerefMut};
14use pocket_ic::{PocketIc, PocketIcBuilder};
15use serde::de::DeserializeOwned;
16
17const INSTALL_CYCLES: u128 = 500 * TC;
18
19#[must_use]
28pub fn pic() -> Pic {
29 PicBuilder::new().with_application_subnet().build()
30}
31
32pub struct PicBuilder(PocketIcBuilder);
42
43#[allow(clippy::new_without_default)]
44impl PicBuilder {
45 #[must_use]
47 pub fn new() -> Self {
48 Self(PocketIcBuilder::new())
49 }
50
51 #[must_use]
53 pub fn with_application_subnet(mut self) -> Self {
54 self.0 = self.0.with_application_subnet();
55 self
56 }
57
58 #[must_use]
60 pub fn with_nns_subnet(mut self) -> Self {
61 self.0 = self.0.with_nns_subnet();
62 self
63 }
64
65 #[must_use]
67 pub fn build(self) -> Pic {
68 Pic(self.0.build())
69 }
70}
71
72#[derive(Deref, DerefMut)]
80pub struct Pic(PocketIc);
81
82impl Pic {
83 pub fn create_and_install_root_canister(&self, wasm: Vec<u8>) -> Result<Principal, Error> {
85 let init_bytes = install_root_args()?;
86
87 Ok(self.create_funded_and_install(wasm, init_bytes))
88 }
89
90 pub fn create_and_install_canister(
94 &self,
95 role: CanisterRole,
96 wasm: Vec<u8>,
97 ) -> Result<Principal, Error> {
98 let init_bytes = install_args(role)?;
99
100 Ok(self.create_funded_and_install(wasm, init_bytes))
101 }
102
103 pub fn create_and_install_canister_with_directories(
108 &self,
109 role: CanisterRole,
110 wasm: Vec<u8>,
111 app_directory: AppDirectoryView,
112 subnet_directory: SubnetDirectoryView,
113 ) -> Result<Principal, Error> {
114 let init_bytes = install_args_with_directories(role, app_directory, subnet_directory)?;
115
116 Ok(self.create_funded_and_install(wasm, init_bytes))
117 }
118
119 fn create_funded_and_install(&self, wasm: Vec<u8>, init_bytes: Vec<u8>) -> Principal {
120 let canister_id = self.create_canister();
121 self.add_cycles(canister_id, INSTALL_CYCLES);
122 self.0.install_canister(canister_id, wasm, init_bytes, None);
123
124 canister_id
125 }
126
127 pub fn update_call<T, A>(
129 &self,
130 canister_id: Principal,
131 method: &str,
132 args: A,
133 ) -> Result<T, Error>
134 where
135 T: CandidType + DeserializeOwned,
136 A: ArgumentEncoder,
137 {
138 let bytes: Vec<u8> = encode_args(args)
139 .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
140 let result = self
141 .0
142 .update_call(canister_id, Principal::anonymous(), method, bytes)
143 .map_err(|err| {
144 Error::internal(format!(
145 "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
146 ))
147 })?;
148
149 decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
150 }
151
152 pub fn update_call_as<T, A>(
154 &self,
155 canister_id: Principal,
156 caller: Principal,
157 method: &str,
158 args: A,
159 ) -> Result<T, Error>
160 where
161 T: CandidType + DeserializeOwned,
162 A: ArgumentEncoder,
163 {
164 let bytes: Vec<u8> = encode_args(args)
165 .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
166 let result = self
167 .0
168 .update_call(canister_id, caller, method, bytes)
169 .map_err(|err| {
170 Error::internal(format!(
171 "pocket_ic update_call failed (canister={canister_id}, method={method}): {err}"
172 ))
173 })?;
174
175 decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
176 }
177
178 pub fn query_call<T, A>(
180 &self,
181 canister_id: 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 .0
193 .query_call(canister_id, Principal::anonymous(), method, bytes)
194 .map_err(|err| {
195 Error::internal(format!(
196 "pocket_ic query_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_as<T, A>(
205 &self,
206 canister_id: Principal,
207 caller: Principal,
208 method: &str,
209 args: A,
210 ) -> Result<T, Error>
211 where
212 T: CandidType + DeserializeOwned,
213 A: ArgumentEncoder,
214 {
215 let bytes: Vec<u8> = encode_args(args)
216 .map_err(|err| Error::internal(format!("encode_args failed: {err}")))?;
217 let result = self
218 .0
219 .query_call(canister_id, caller, method, bytes)
220 .map_err(|err| {
221 Error::internal(format!(
222 "pocket_ic query_call failed (canister={canister_id}, method={method}): {err}"
223 ))
224 })?;
225
226 decode_one(&result).map_err(|err| Error::internal(format!("decode_one failed: {err}")))
227 }
228
229 pub fn tick_n(&self, times: usize) {
230 for _ in 0..times {
231 self.tick();
232 }
233 }
234}
235
236fn install_args(role: CanisterRole) -> Result<Vec<u8>, Error> {
250 if role.is_root() {
251 install_root_args()
252 } else {
253 let env = EnvView {
256 prime_root_pid: None,
257 subnet_role: None,
258 subnet_pid: None,
259 root_pid: None,
260 canister_role: Some(role),
261 parent_pid: None,
262 };
263
264 let payload = CanisterInitPayload {
267 env,
268 app_directory: AppDirectoryView(Vec::new()),
269 subnet_directory: SubnetDirectoryView(Vec::new()),
270 };
271
272 encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
273 .map_err(|err| Error::internal(format!("encode_args failed: {err}")))
274 }
275}
276
277fn install_root_args() -> Result<Vec<u8>, Error> {
278 encode_one(SubnetIdentity::Manual)
279 .map_err(|err| Error::internal(format!("encode_one failed: {err}")))
280}
281
282fn install_args_with_directories(
283 role: CanisterRole,
284 app_directory: AppDirectoryView,
285 subnet_directory: SubnetDirectoryView,
286) -> Result<Vec<u8>, Error> {
287 if role.is_root() {
288 encode_one(SubnetIdentity::Manual)
291 .map_err(|err| Error::internal(format!("encode_one failed: {err}")))
292 } else {
293 let env = EnvView {
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 {
304 env,
305 app_directory,
306 subnet_directory,
307 };
308
309 encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
310 .map_err(|err| Error::internal(format!("encode_args failed: {err}")))
311 }
312}