1use candid::{CandidType, Principal, decode_one, encode_args, encode_one, utils::ArgumentEncoder};
2use canic::{
3 Error,
4 core::{
5 ids::{CanisterRole, SubnetRole},
6 ops::storage::{
7 CanisterInitPayload, directory::DirectoryView, env::EnvData, topology::SubnetIdentity,
8 },
9 },
10};
11use derive_more::{Deref, DerefMut};
12use pocket_ic::{PocketIc, PocketIcBuilder};
13use serde::de::DeserializeOwned;
14use std::sync::OnceLock;
15
16static PIC: OnceLock<Pic> = OnceLock::new();
33
34#[must_use]
40pub fn pic() -> &'static Pic {
41 PIC.get_or_init(|| PicBuilder::new().with_application_subnet().build())
42}
43
44pub struct PicBuilder(PocketIcBuilder);
54
55#[allow(clippy::new_without_default)]
56impl PicBuilder {
57 #[must_use]
59 pub fn new() -> Self {
60 Self(PocketIcBuilder::new())
61 }
62
63 #[must_use]
65 pub fn with_application_subnet(mut self) -> Self {
66 self.0 = self.0.with_application_subnet();
67 self
68 }
69
70 #[must_use]
72 pub fn with_nns_subnet(mut self) -> Self {
73 self.0 = self.0.with_nns_subnet();
74 self
75 }
76
77 #[must_use]
79 pub fn build(self) -> Pic {
80 Pic(self.0.build())
81 }
82}
83
84#[derive(Deref, DerefMut)]
92pub struct Pic(PocketIc);
93
94impl Pic {
95 pub fn create_and_install_canister(
99 &self,
100 role: CanisterRole,
101 wasm: Vec<u8>,
102 ) -> Result<Principal, Error> {
103 let canister_id = self.create_canister();
105 self.add_cycles(canister_id, 1_000_000_000_000);
106
107 let init_bytes = install_args(role)?;
109 self.0.install_canister(canister_id, wasm, init_bytes, None);
110
111 Ok(canister_id)
112 }
113
114 pub fn create_and_install_canister_with_directories(
119 &self,
120 role: CanisterRole,
121 wasm: Vec<u8>,
122 app_directory: DirectoryView,
123 subnet_directory: DirectoryView,
124 ) -> Result<Principal, Error> {
125 let canister_id = self.create_canister();
126 self.add_cycles(canister_id, 1_000_000_000_000);
127
128 let init_bytes = install_args_with_directories(role, app_directory, subnet_directory)?;
129 self.0.install_canister(canister_id, wasm, init_bytes, None);
130
131 Ok(canister_id)
132 }
133
134 pub fn update_call<T, A>(
136 &self,
137 canister_id: Principal,
138 method: &str,
139 args: A,
140 ) -> Result<T, Error>
141 where
142 T: CandidType + DeserializeOwned,
143 A: ArgumentEncoder,
144 {
145 let bytes: Vec<u8> = encode_args(args)?;
146 let result = self
147 .0
148 .update_call(canister_id, Principal::anonymous(), method, bytes)
149 .map_err(|e| Error::test(e.to_string()))?;
150
151 decode_one(&result).map_err(Into::into)
152 }
153
154 pub fn query_call<T, A>(
156 &self,
157 canister_id: Principal,
158 method: &str,
159 args: A,
160 ) -> Result<T, Error>
161 where
162 T: CandidType + DeserializeOwned,
163 A: ArgumentEncoder,
164 {
165 let bytes: Vec<u8> = encode_args(args)?;
166 let result = self
167 .0
168 .query_call(canister_id, Principal::anonymous(), method, bytes)
169 .map_err(|e| Error::test(e.to_string()))?;
170
171 decode_one(&result).map_err(Into::into)
172 }
173}
174
175fn install_args(role: CanisterRole) -> Result<Vec<u8>, Error> {
189 let args = if role.is_root() {
190 let subnet_pid = Principal::from_slice(&[0xAA; 29]);
192 encode_one(SubnetIdentity::Manual(subnet_pid))
193 } else {
194 let root_pid = Principal::from_slice(&[0xBB; 29]);
196 let subnet_pid = Principal::from_slice(&[0xAA; 29]);
197 let env = EnvData {
198 prime_root_pid: Some(root_pid),
199 subnet_role: Some(SubnetRole::PRIME),
200 subnet_pid: Some(subnet_pid),
201 root_pid: Some(root_pid),
202 canister_role: Some(role),
203 parent_pid: Some(root_pid),
204 };
205 let payload = CanisterInitPayload::new(env, Vec::new(), Vec::new());
208 encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
209 }?;
210
211 Ok(args)
212}
213
214fn install_args_with_directories(
215 role: CanisterRole,
216 app_directory: DirectoryView,
217 subnet_directory: DirectoryView,
218) -> Result<Vec<u8>, Error> {
219 let args = if role.is_root() {
220 let subnet_pid = Principal::from_slice(&[0xAA; 29]);
221 encode_one(SubnetIdentity::Manual(subnet_pid))
222 } else {
223 let root_pid = Principal::from_slice(&[0xBB; 29]);
224 let subnet_pid = Principal::from_slice(&[0xAA; 29]);
225 let env = EnvData {
226 prime_root_pid: Some(root_pid),
227 subnet_role: Some(SubnetRole::PRIME),
228 subnet_pid: Some(subnet_pid),
229 root_pid: Some(root_pid),
230 canister_role: Some(role),
231 parent_pid: Some(root_pid),
232 };
233 let payload = CanisterInitPayload::new(env, app_directory, subnet_directory);
234 encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
235 }?;
236
237 Ok(args)
238}