1use candid::{CandidType, Principal, decode_one, encode_args, encode_one, utils::ArgumentEncoder};
2use canic::{
3 Error,
4 core::{
5 ids::CanisterRole,
6 ops::storage::{CanisterInitPayload, topology::SubnetIdentity},
7 },
8};
9use derive_more::{Deref, DerefMut};
10use pocket_ic::{PocketIc, PocketIcBuilder};
11use serde::de::DeserializeOwned;
12
13pub struct PicBuilder(PocketIcBuilder);
18
19#[allow(clippy::new_without_default)]
20impl PicBuilder {
21 #[must_use]
23 pub fn new() -> Self {
24 Self(PocketIcBuilder::new())
25 }
26
27 #[must_use]
28 pub fn with_application_subnet(mut self) -> Self {
29 self.0 = self.0.with_application_subnet();
30 self
31 }
32
33 #[must_use]
34 pub fn with_nns_subnet(mut self) -> Self {
35 self.0 = self.0.with_nns_subnet();
36 self
37 }
38
39 #[must_use]
41 pub fn build(self) -> Pic {
42 Pic(self.0.build())
43 }
44}
45
46#[derive(Deref, DerefMut)]
51pub struct Pic(PocketIc);
52
53impl Pic {
54 pub fn create_and_install_canister(
56 &self,
57 role: CanisterRole,
58 wasm: Vec<u8>,
59 ) -> Result<Principal, Error> {
60 let canister_id = self.create_canister();
62 self.add_cycles(canister_id, 1_000_000_000_000);
63
64 let init_bytes = install_args(role)?;
66 self.0.install_canister(canister_id, wasm, init_bytes, None);
67
68 Ok(canister_id)
69 }
70
71 pub fn update_call<T, A>(
73 &self,
74 canister_id: Principal,
75 method: &str,
76 args: A,
77 ) -> Result<T, Error>
78 where
79 T: CandidType + DeserializeOwned,
80 A: ArgumentEncoder,
81 {
82 let bytes: Vec<u8> = encode_args(args)?;
83 let result = self
84 .0
85 .update_call(canister_id, Principal::anonymous(), method, bytes)
86 .map_err(|e| Error::test(e.to_string()))?;
87
88 decode_one(&result).map_err(Into::into)
89 }
90
91 pub fn query_call<T, A>(
93 &self,
94 canister_id: Principal,
95 method: &str,
96 args: A,
97 ) -> Result<T, Error>
98 where
99 T: CandidType + DeserializeOwned,
100 A: ArgumentEncoder,
101 {
102 let bytes: Vec<u8> = encode_args(args)?;
103 let result = self
104 .0
105 .query_call(canister_id, Principal::anonymous(), method, bytes)
106 .map_err(|e| Error::test(e.to_string()))?;
107
108 decode_one(&result).map_err(Into::into)
109 }
110}
111
112fn install_args(role: CanisterRole) -> Result<Vec<u8>, Error> {
116 let args = if role.is_root() {
117 let subnet_pid = Principal::from_slice(&[0xAA; 29]);
119 encode_one(SubnetIdentity::Manual(subnet_pid))
120 } else {
121 let payload = CanisterInitPayload::empty();
122 encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
123 }?;
124
125 Ok(args)
126}