Skip to main content

ic_testkit/pic/
calls.rs

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
22/// Typed Candid calls with contextual encoding, rejection, and decoding errors.
23pub trait CandidCallExt {
24    /// Encode and execute an anonymous update call, then decode its result.
25    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    /// Execute [`update_candid`](Self::update_candid), panicking on harness errors.
36    #[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    /// Encode and execute an update call as an explicit caller.
43    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    /// Execute [`update_candid_as`](Self::update_candid_as), panicking on harness errors.
55    #[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    /// Encode and execute an anonymous query call, then decode its result.
68    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    /// Execute [`query_candid`](Self::query_candid), panicking on harness errors.
79    #[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    /// Encode and execute a query call as an explicit caller.
86    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    /// Execute [`query_candid_as`](Self::query_candid_as), panicking on harness errors.
98    #[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    /// Generic update call helper (serializes args + decodes result).
113    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    /// Generic update call helper that panics on rejection or Candid codec failure.
127    ///
128    /// This does not unwrap application-level results. For example,
129    /// `update_candid_or_panic::<Result<T, E>, _>(...)` returns `Result<T, E>`.
130    #[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    /// Generic update call helper with an explicit caller principal.
141    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    /// Generic update call helper with an explicit caller principal that panics
167    /// on rejection or Candid codec failure.
168    ///
169    /// This does not unwrap application-level results. For example,
170    /// `update_candid_as_or_panic::<Result<T, E>, _>(...)` returns `Result<T, E>`.
171    #[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    /// Generic query call helper.
188    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    /// Generic query call helper that panics on rejection or Candid codec failure.
202    ///
203    /// This does not unwrap application-level results. For example,
204    /// `query_candid_or_panic::<Result<T, E>, _>(...)` returns `Result<T, E>`.
205    #[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    /// Generic query call helper with an explicit caller principal.
216    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    /// Generic query call helper with an explicit caller principal that panics
242    /// on rejection or Candid codec failure.
243    ///
244    /// This does not unwrap application-level results. For example,
245    /// `query_candid_as_or_panic::<Result<T, E>, _>(...)` returns `Result<T, E>`.
246    #[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}