1use candid::{CandidType, Principal, decode_one, encode_args, utils::ArgumentEncoder};
2use pocket_ic::{PocketIc, RejectResponse};
3use serde::de::DeserializeOwned;
4use std::panic::{AssertUnwindSafe, catch_unwind, resume_unwind};
5
6use super::{CandidCallContext, CandidCallError, startup};
7
8#[derive(Clone, Copy)]
9struct CallContext<'a> {
10 operation: &'static str,
11 canister_id: Principal,
12 caller: Principal,
13 method: &'a str,
14}
15
16impl CallContext<'_> {
17 fn to_error_context(self) -> CandidCallContext {
18 CandidCallContext::new(self.operation, self.canister_id, self.caller, self.method)
19 }
20}
21
22pub trait CandidCallExt {
24 fn update_candid<T, A>(
26 &self,
27 canister_id: Principal,
28 method: &str,
29 args: A,
30 ) -> Result<T, CandidCallError>
31 where
32 T: CandidType + DeserializeOwned,
33 A: ArgumentEncoder;
34
35 #[track_caller]
37 fn update_candid_or_panic<T, A>(&self, canister_id: Principal, method: &str, args: A) -> T
38 where
39 T: CandidType + DeserializeOwned,
40 A: ArgumentEncoder;
41
42 fn update_candid_as<T, A>(
44 &self,
45 canister_id: Principal,
46 caller: Principal,
47 method: &str,
48 args: A,
49 ) -> Result<T, CandidCallError>
50 where
51 T: CandidType + DeserializeOwned,
52 A: ArgumentEncoder;
53
54 #[track_caller]
56 fn update_candid_as_or_panic<T, A>(
57 &self,
58 canister_id: Principal,
59 caller: Principal,
60 method: &str,
61 args: A,
62 ) -> T
63 where
64 T: CandidType + DeserializeOwned,
65 A: ArgumentEncoder;
66
67 fn query_candid<T, A>(
69 &self,
70 canister_id: Principal,
71 method: &str,
72 args: A,
73 ) -> Result<T, CandidCallError>
74 where
75 T: CandidType + DeserializeOwned,
76 A: ArgumentEncoder;
77
78 #[track_caller]
80 fn query_candid_or_panic<T, A>(&self, canister_id: Principal, method: &str, args: A) -> T
81 where
82 T: CandidType + DeserializeOwned,
83 A: ArgumentEncoder;
84
85 fn query_candid_as<T, A>(
87 &self,
88 canister_id: Principal,
89 caller: Principal,
90 method: &str,
91 args: A,
92 ) -> Result<T, CandidCallError>
93 where
94 T: CandidType + DeserializeOwned,
95 A: ArgumentEncoder;
96
97 #[track_caller]
99 fn query_candid_as_or_panic<T, A>(
100 &self,
101 canister_id: Principal,
102 caller: Principal,
103 method: &str,
104 args: A,
105 ) -> T
106 where
107 T: CandidType + DeserializeOwned,
108 A: ArgumentEncoder;
109}
110
111impl CandidCallExt for PocketIc {
112 fn update_candid<T, A>(
114 &self,
115 canister_id: Principal,
116 method: &str,
117 args: A,
118 ) -> Result<T, CandidCallError>
119 where
120 T: CandidType + DeserializeOwned,
121 A: ArgumentEncoder,
122 {
123 self.update_candid_as(canister_id, Principal::anonymous(), method, args)
124 }
125
126 #[track_caller]
131 fn update_candid_or_panic<T, A>(&self, canister_id: Principal, method: &str, args: A) -> T
132 where
133 T: CandidType + DeserializeOwned,
134 A: ArgumentEncoder,
135 {
136 self.update_candid(canister_id, method, args)
137 .unwrap_or_else(|err| panic!("{err}"))
138 }
139
140 fn update_candid_as<T, A>(
142 &self,
143 canister_id: Principal,
144 caller: Principal,
145 method: &str,
146 args: A,
147 ) -> Result<T, CandidCallError>
148 where
149 T: CandidType + DeserializeOwned,
150 A: ArgumentEncoder,
151 {
152 let context = CallContext {
153 operation: "update_call",
154 canister_id,
155 caller,
156 method,
157 };
158 let bytes = encode_call_args(args, context)?;
159 let result = run_raw_call(context, || {
160 Self::update_call(self, canister_id, caller, method, bytes)
161 })?;
162
163 decode_call_result(&result, context)
164 }
165
166 #[track_caller]
172 fn update_candid_as_or_panic<T, A>(
173 &self,
174 canister_id: Principal,
175 caller: Principal,
176 method: &str,
177 args: A,
178 ) -> T
179 where
180 T: CandidType + DeserializeOwned,
181 A: ArgumentEncoder,
182 {
183 self.update_candid_as(canister_id, caller, method, args)
184 .unwrap_or_else(|err| panic!("{err}"))
185 }
186
187 fn query_candid<T, A>(
189 &self,
190 canister_id: Principal,
191 method: &str,
192 args: A,
193 ) -> Result<T, CandidCallError>
194 where
195 T: CandidType + DeserializeOwned,
196 A: ArgumentEncoder,
197 {
198 self.query_candid_as(canister_id, Principal::anonymous(), method, args)
199 }
200
201 #[track_caller]
206 fn query_candid_or_panic<T, A>(&self, canister_id: Principal, method: &str, args: A) -> T
207 where
208 T: CandidType + DeserializeOwned,
209 A: ArgumentEncoder,
210 {
211 self.query_candid(canister_id, method, args)
212 .unwrap_or_else(|err| panic!("{err}"))
213 }
214
215 fn query_candid_as<T, A>(
217 &self,
218 canister_id: Principal,
219 caller: Principal,
220 method: &str,
221 args: A,
222 ) -> Result<T, CandidCallError>
223 where
224 T: CandidType + DeserializeOwned,
225 A: ArgumentEncoder,
226 {
227 let context = CallContext {
228 operation: "query_call",
229 canister_id,
230 caller,
231 method,
232 };
233 let bytes = encode_call_args(args, context)?;
234 let result = run_raw_call(context, || {
235 Self::query_call(self, canister_id, caller, method, bytes)
236 })?;
237
238 decode_call_result(&result, context)
239 }
240
241 #[track_caller]
247 fn query_candid_as_or_panic<T, A>(
248 &self,
249 canister_id: Principal,
250 caller: Principal,
251 method: &str,
252 args: A,
253 ) -> T
254 where
255 T: CandidType + DeserializeOwned,
256 A: ArgumentEncoder,
257 {
258 self.query_candid_as(canister_id, caller, method, args)
259 .unwrap_or_else(|err| panic!("{err}"))
260 }
261}
262
263fn encode_call_args<A>(args: A, context: CallContext<'_>) -> Result<Vec<u8>, CandidCallError>
264where
265 A: ArgumentEncoder,
266{
267 encode_args(args).map_err(|err| CandidCallError::encode(context.to_error_context(), err))
268}
269
270fn decode_call_result<T>(result: &[u8], context: CallContext<'_>) -> Result<T, CandidCallError>
271where
272 T: CandidType + DeserializeOwned,
273{
274 decode_one(result)
275 .map_err(|err| CandidCallError::decode(context.to_error_context(), result.len(), err))
276}
277
278fn run_raw_call<F>(context: CallContext<'_>, call: F) -> Result<Vec<u8>, CandidCallError>
279where
280 F: FnOnce() -> Result<Vec<u8>, RejectResponse>,
281{
282 match catch_unwind(AssertUnwindSafe(call)) {
283 Ok(Ok(bytes)) => Ok(bytes),
284 Ok(Err(response)) => Err(CandidCallError::canister_reject(
285 context.to_error_context(),
286 response,
287 )),
288 Err(payload) if startup::panic_is_dead_instance_transport(payload.as_ref()) => {
289 Err(CandidCallError::transport(
290 context.to_error_context(),
291 startup::panic_payload_to_string(payload.as_ref()),
292 ))
293 }
294 Err(payload) => resume_unwind(payload),
295 }
296}
297
298#[cfg(test)]
299mod tests {
300 use candid::Principal;
301
302 use crate::pic::CandidCallErrorKind;
303
304 use super::{CallContext, decode_call_result, run_raw_call};
305
306 #[test]
307 fn decode_error_includes_call_context() {
308 let context = CallContext {
309 operation: "query_call",
310 canister_id: Principal::anonymous(),
311 caller: Principal::management_canister(),
312 method: "get",
313 };
314
315 let err = decode_call_result::<u64>(&[0xde, 0xad], context).expect_err("decode fails");
316
317 assert!(err.message().contains("candid decode_one failed"));
318 assert!(err.message().contains("operation=query_call"));
319 assert!(err.message().contains("method=get"));
320 assert!(err.message().contains("bytes=2"));
321 assert_eq!(err.kind(), CandidCallErrorKind::Decode);
322 assert_eq!(err.context().expect("decode error context").method(), "get");
323 }
324
325 #[test]
326 fn dead_instance_panic_is_classified_as_transport() {
327 let context = CallContext {
328 operation: "query_call",
329 canister_id: Principal::anonymous(),
330 caller: Principal::management_canister(),
331 method: "get",
332 };
333 let error = run_raw_call(context, || -> Result<Vec<u8>, pocket_ic::RejectResponse> {
334 panic!("transport failed: ConnectionRefused");
335 })
336 .unwrap_err();
337
338 assert_eq!(error.kind(), crate::pic::CandidCallErrorKind::Transport);
339 assert!(error.reject_response().is_none());
340 }
341}