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;
12use std::sync::OnceLock;
13
14static PIC: OnceLock<Pic> = OnceLock::new();
19
20pub fn pic() -> &'static Pic {
21 PIC.get_or_init(|| PicBuilder::new().with_application_subnet().build())
22}
23
24pub struct PicBuilder(PocketIcBuilder);
29
30#[allow(clippy::new_without_default)]
31impl PicBuilder {
32 #[must_use]
34 pub fn new() -> Self {
35 Self(PocketIcBuilder::new())
36 }
37
38 #[must_use]
39 pub fn with_application_subnet(mut self) -> Self {
40 self.0 = self.0.with_application_subnet();
41 self
42 }
43
44 #[must_use]
45 pub fn with_nns_subnet(mut self) -> Self {
46 self.0 = self.0.with_nns_subnet();
47 self
48 }
49
50 #[must_use]
52 pub fn build(self) -> Pic {
53 Pic(self.0.build())
54 }
55}
56
57#[derive(Deref, DerefMut)]
62pub struct Pic(PocketIc);
63
64impl Pic {
65 pub fn create_and_install_canister(
67 &self,
68 role: CanisterRole,
69 wasm: Vec<u8>,
70 ) -> Result<Principal, Error> {
71 let canister_id = self.create_canister();
73 self.add_cycles(canister_id, 1_000_000_000_000);
74
75 let init_bytes = install_args(role)?;
77 self.0.install_canister(canister_id, wasm, init_bytes, None);
78
79 Ok(canister_id)
80 }
81
82 pub fn update_call<T, A>(
84 &self,
85 canister_id: Principal,
86 method: &str,
87 args: A,
88 ) -> Result<T, Error>
89 where
90 T: CandidType + DeserializeOwned,
91 A: ArgumentEncoder,
92 {
93 let bytes: Vec<u8> = encode_args(args)?;
94 let result = self
95 .0
96 .update_call(canister_id, Principal::anonymous(), method, bytes)
97 .map_err(|e| Error::test(e.to_string()))?;
98
99 decode_one(&result).map_err(Into::into)
100 }
101
102 pub fn query_call<T, A>(
104 &self,
105 canister_id: Principal,
106 method: &str,
107 args: A,
108 ) -> Result<T, Error>
109 where
110 T: CandidType + DeserializeOwned,
111 A: ArgumentEncoder,
112 {
113 let bytes: Vec<u8> = encode_args(args)?;
114 let result = self
115 .0
116 .query_call(canister_id, Principal::anonymous(), method, bytes)
117 .map_err(|e| Error::test(e.to_string()))?;
118
119 decode_one(&result).map_err(Into::into)
120 }
121}
122
123fn install_args(role: CanisterRole) -> Result<Vec<u8>, Error> {
127 let args = if role.is_root() {
128 let subnet_pid = Principal::from_slice(&[0xAA; 29]);
130 encode_one(SubnetIdentity::Manual(subnet_pid))
131 } else {
132 let payload = CanisterInitPayload::empty();
133 encode_args::<(CanisterInitPayload, Option<Vec<u8>>)>((payload, None))
134 }?;
135
136 Ok(args)
137}