Skip to main content

ic_testkit/pic/
calls.rs

1use candid::{
2    CandidType, Principal, decode_args, decode_one, encode_args,
3    utils::{ArgumentDecoder, ArgumentEncoder},
4};
5use pocket_ic::{PocketIc, RejectResponse};
6use serde::de::DeserializeOwned;
7use std::panic::{AssertUnwindSafe, catch_unwind, resume_unwind};
8
9use super::{CandidCallContext, CandidCallError, transport};
10
11#[derive(Clone, Copy)]
12struct CallContext<'a> {
13    operation: &'static str,
14    canister_id: Principal,
15    caller: Principal,
16    method: &'a str,
17}
18
19impl CallContext<'_> {
20    fn to_error_context(self) -> CandidCallContext {
21        CandidCallContext::new(self.operation, self.canister_id, self.caller, self.method)
22    }
23}
24
25/// Typed Candid calls with contextual encoding, rejection, and decoding errors.
26pub trait CandidCallExt {
27    /// Encode and execute an anonymous update call, then decode its result.
28    fn update_candid<T, A>(
29        &self,
30        canister_id: Principal,
31        method: &str,
32        args: A,
33    ) -> Result<T, CandidCallError>
34    where
35        T: CandidType + DeserializeOwned,
36        A: ArgumentEncoder,
37    {
38        self.update_candid_as(canister_id, Principal::anonymous(), method, args)
39    }
40
41    /// Encode and execute an update call as an explicit caller.
42    fn update_candid_as<T, A>(
43        &self,
44        canister_id: Principal,
45        caller: Principal,
46        method: &str,
47        args: A,
48    ) -> Result<T, CandidCallError>
49    where
50        T: CandidType + DeserializeOwned,
51        A: ArgumentEncoder;
52
53    /// Encode and execute an anonymous query call, then decode its result.
54    fn query_candid<T, A>(
55        &self,
56        canister_id: Principal,
57        method: &str,
58        args: A,
59    ) -> Result<T, CandidCallError>
60    where
61        T: CandidType + DeserializeOwned,
62        A: ArgumentEncoder,
63    {
64        self.query_candid_as(canister_id, Principal::anonymous(), method, args)
65    }
66
67    /// Encode and execute a query call as an explicit caller.
68    fn query_candid_as<T, A>(
69        &self,
70        canister_id: Principal,
71        caller: Principal,
72        method: &str,
73        args: A,
74    ) -> Result<T, CandidCallError>
75    where
76        T: CandidType + DeserializeOwned,
77        A: ArgumentEncoder;
78}
79
80impl CandidCallExt for PocketIc {
81    /// Generic update call helper with an explicit caller principal.
82    fn update_candid_as<T, A>(
83        &self,
84        canister_id: Principal,
85        caller: Principal,
86        method: &str,
87        args: A,
88    ) -> Result<T, CandidCallError>
89    where
90        T: CandidType + DeserializeOwned,
91        A: ArgumentEncoder,
92    {
93        let context = CallContext {
94            operation: "update_call",
95            canister_id,
96            caller,
97            method,
98        };
99        call_candid(
100            context,
101            args,
102            |bytes| Self::update_call(self, canister_id, caller, method, bytes),
103            |bytes| decode_one(bytes),
104        )
105    }
106
107    /// Generic query call helper with an explicit caller principal.
108    fn query_candid_as<T, A>(
109        &self,
110        canister_id: Principal,
111        caller: Principal,
112        method: &str,
113        args: A,
114    ) -> Result<T, CandidCallError>
115    where
116        T: CandidType + DeserializeOwned,
117        A: ArgumentEncoder,
118    {
119        let context = CallContext {
120            operation: "query_call",
121            canister_id,
122            caller,
123            method,
124        };
125        call_candid(
126            context,
127            args,
128            |bytes| Self::query_call(self, canister_id, caller, method, bytes),
129            |bytes| decode_one(bytes),
130        )
131    }
132}
133
134/// Execute an update as `caller`, decoding the top-level Candid reply arguments.
135///
136/// Unlike [`CandidCallExt::update_candid_as`], `(u64, bool)` means two reply
137/// arguments rather than one tuple-valued reply. Use `()` for no replies or
138/// `(T,)` for one reply. Pass [`Principal::anonymous()`] for an anonymous caller.
139/// Application-level `Result<T, E>` values are returned without unwrapping.
140pub fn update_candid_args<T, A>(
141    pic: &PocketIc,
142    canister_id: Principal,
143    caller: Principal,
144    method: &str,
145    args: A,
146) -> Result<T, CandidCallError>
147where
148    T: for<'de> ArgumentDecoder<'de>,
149    A: ArgumentEncoder,
150{
151    let context = CallContext {
152        operation: "update_call",
153        canister_id,
154        caller,
155        method,
156    };
157    call_candid(
158        context,
159        args,
160        |bytes| pic.update_call(canister_id, caller, method, bytes),
161        |bytes| decode_args(bytes),
162    )
163}
164
165/// Execute a query as `caller`, decoding the top-level Candid reply arguments.
166///
167/// Unlike [`CandidCallExt::query_candid_as`], `(u64, bool)` means two reply
168/// arguments rather than one tuple-valued reply. Use `()` for no replies or
169/// `(T,)` for one reply. Pass [`Principal::anonymous()`] for an anonymous caller.
170/// Application-level `Result<T, E>` values are returned without unwrapping.
171pub fn query_candid_args<T, A>(
172    pic: &PocketIc,
173    canister_id: Principal,
174    caller: Principal,
175    method: &str,
176    args: A,
177) -> Result<T, CandidCallError>
178where
179    T: for<'de> ArgumentDecoder<'de>,
180    A: ArgumentEncoder,
181{
182    let context = CallContext {
183        operation: "query_call",
184        canister_id,
185        caller,
186        method,
187    };
188    call_candid(
189        context,
190        args,
191        |bytes| pic.query_call(canister_id, caller, method, bytes),
192        |bytes| decode_args(bytes),
193    )
194}
195
196fn call_candid<T, A, F, D>(
197    context: CallContext<'_>,
198    args: A,
199    call: F,
200    decode: D,
201) -> Result<T, CandidCallError>
202where
203    A: ArgumentEncoder,
204    F: FnOnce(Vec<u8>) -> Result<Vec<u8>, RejectResponse>,
205    D: FnOnce(&[u8]) -> Result<T, candid::Error>,
206{
207    let bytes = encode_args(args)
208        .map_err(|err| CandidCallError::encode(context.to_error_context(), err))?;
209    let result = run_raw_call(context, || call(bytes))?;
210    decode(&result)
211        .map_err(|err| CandidCallError::decode(context.to_error_context(), result.len(), err))
212}
213
214fn run_raw_call<F>(context: CallContext<'_>, call: F) -> Result<Vec<u8>, CandidCallError>
215where
216    F: FnOnce() -> Result<Vec<u8>, RejectResponse>,
217{
218    match catch_unwind(AssertUnwindSafe(call)) {
219        Ok(Ok(bytes)) => Ok(bytes),
220        Ok(Err(response)) => Err(CandidCallError::canister_reject(
221            context.to_error_context(),
222            response,
223        )),
224        Err(payload) if transport::panic_is_dead_instance_transport(payload.as_ref()) => {
225            Err(CandidCallError::transport(
226                context.to_error_context(),
227                transport::panic_payload_to_string(payload.as_ref()),
228            ))
229        }
230        Err(payload) => resume_unwind(payload),
231    }
232}
233
234#[cfg(test)]
235mod tests {
236    use candid::Principal;
237
238    use crate::pic::CandidCallErrorKind;
239
240    use super::{CallContext, call_candid, run_raw_call};
241
242    #[test]
243    fn decode_error_includes_call_context() {
244        let context = CallContext {
245            operation: "query_call",
246            canister_id: Principal::anonymous(),
247            caller: Principal::management_canister(),
248            method: "get",
249        };
250
251        let err = call_candid(
252            context,
253            (),
254            |_| Ok(vec![0xde, 0xad]),
255            |bytes| candid::decode_one::<u64>(bytes),
256        )
257        .expect_err("decode fails");
258
259        assert!(err.message.contains("candid decode failed"));
260        assert!(err.message.contains("operation=query_call"));
261        assert!(err.message.contains("method=get"));
262        assert!(err.message.contains("bytes=2"));
263        assert_eq!(err.kind, CandidCallErrorKind::Decode);
264        assert_eq!(
265            err.context.as_deref().expect("decode error context").method,
266            "get"
267        );
268    }
269
270    #[test]
271    fn reply_arguments_preserve_top_level_arity_and_wide_values() {
272        let context = CallContext {
273            operation: "update_call",
274            canister_id: Principal::anonymous(),
275            caller: Principal::management_canister(),
276            method: "measure",
277        };
278        let values = (u128::MAX, u128::from(u64::MAX) + 1);
279        let bytes = candid::encode_args(values).unwrap();
280        assert_eq!(candid::decode_args::<(u128, u128)>(&bytes).unwrap(), values);
281        assert!(candid::decode_one::<(u128, u128)>(&bytes).is_err());
282        let tuple = candid::encode_one(values).unwrap();
283        assert!(candid::decode_args::<(u128, u128)>(&tuple).is_err());
284        assert_eq!(candid::decode_one::<(u128, u128)>(&tuple).unwrap(), values);
285        let no_replies = candid::encode_args(()).unwrap();
286        candid::decode_args::<()>(&no_replies).unwrap();
287        let one_reply = candid::encode_args((7_u64,)).unwrap();
288        assert_eq!(candid::decode_args::<(u64,)>(&one_reply).unwrap(), (7,));
289        let result: (u64, bool, String) = call_candid(
290            context,
291            (42_u64, false),
292            |args| {
293                assert_eq!(
294                    candid::decode_args::<(u64, bool)>(&args).unwrap(),
295                    (42, false)
296                );
297                Ok(candid::encode_args((9_u64, true, "ready")).unwrap())
298            },
299            |bytes| candid::decode_args(bytes),
300        )
301        .unwrap();
302        assert_eq!(result, (9, true, "ready".into()));
303    }
304
305    #[test]
306    fn invalid_reply_arguments_keep_decode_context() {
307        let context = CallContext {
308            operation: "query_call",
309            canister_id: Principal::anonymous(),
310            caller: Principal::management_canister(),
311            method: "measure",
312        };
313        for bytes in [vec![0xde, 0xad], candid::encode_args((7_u64,)).unwrap()] {
314            let length = bytes.len();
315            let err = call_candid(
316                context,
317                (),
318                |_| Ok(bytes),
319                |bytes| candid::decode_args::<(u64, bool)>(bytes),
320            )
321            .unwrap_err();
322            assert_eq!(err.kind, CandidCallErrorKind::Decode);
323            let observed = err.context.as_deref().unwrap();
324            assert_eq!(observed.operation, "query_call");
325            assert_eq!(observed.canister_id, context.canister_id);
326            assert_eq!(observed.caller, context.caller);
327            assert_eq!(observed.method, "measure");
328            assert!(err.message.contains(&format!("bytes={length}")));
329        }
330    }
331
332    #[test]
333    fn dead_instance_panic_is_classified_as_transport() {
334        let context = CallContext {
335            operation: "query_call",
336            canister_id: Principal::anonymous(),
337            caller: Principal::management_canister(),
338            method: "get",
339        };
340        let error = run_raw_call(context, || -> Result<Vec<u8>, pocket_ic::RejectResponse> {
341            panic!("HTTP failure: reqwest::Error {{ kind: Request, url: \"http://127.0.0.1:1234/instances/0/update/execute_ingress_message\", source: ConnectError(\"tcp connect error\", Os {{ code: 111, kind: ConnectionRefused, message: \"Connection refused\" }}) }}");
342        })
343        .unwrap_err();
344
345        assert_eq!(error.kind, crate::pic::CandidCallErrorKind::Transport);
346        assert!(error.reject_response.is_none());
347    }
348
349    #[test]
350    fn unrelated_panics_resume_with_the_original_payload() {
351        for message in [
352            "application worker channel closed",
353            "fixture expected ConnectionRefused but observed another result",
354        ] {
355            let context = CallContext {
356                operation: "query_call",
357                canister_id: Principal::anonymous(),
358                caller: Principal::management_canister(),
359                method: "get",
360            };
361            let payload = message.to_owned();
362            let original_allocation = payload.as_ptr() as usize;
363            let result = std::panic::catch_unwind(|| {
364                run_raw_call(context, || -> Result<Vec<u8>, pocket_ic::RejectResponse> {
365                    std::panic::panic_any(payload);
366                })
367            });
368            let resumed = result
369                .expect_err("application panic must resume")
370                .downcast::<String>()
371                .expect("original panic type must be preserved");
372            assert_eq!(resumed.as_str(), message);
373            assert_eq!(resumed.as_ptr() as usize, original_allocation);
374        }
375    }
376}