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).map_err(|err| CandidCallError::Encode {
208        context: Box::new(context.to_error_context()),
209        message: err.to_string(),
210    })?;
211    let result = run_raw_call(context, || call(bytes))?;
212    decode(&result).map_err(|err| CandidCallError::Decode {
213        context: Box::new(context.to_error_context()),
214        bytes: result.len(),
215        message: err.to_string(),
216    })
217}
218
219fn run_raw_call<F>(context: CallContext<'_>, call: F) -> Result<Vec<u8>, CandidCallError>
220where
221    F: FnOnce() -> Result<Vec<u8>, RejectResponse>,
222{
223    match catch_unwind(AssertUnwindSafe(call)) {
224        Ok(Ok(bytes)) => Ok(bytes),
225        Ok(Err(response)) => Err(CandidCallError::CanisterReject {
226            context: Box::new(context.to_error_context()),
227            response: Box::new(response),
228        }),
229        Err(payload) if transport::panic_is_dead_instance_transport(payload.as_ref()) => {
230            Err(CandidCallError::Transport {
231                context: Box::new(context.to_error_context()),
232                message: transport::panic_payload_to_string(payload.as_ref()),
233            })
234        }
235        Err(payload) => resume_unwind(payload),
236    }
237}
238
239#[cfg(test)]
240mod tests {
241    use candid::Principal;
242
243    use super::{CallContext, CandidCallError, call_candid, run_raw_call};
244
245    #[test]
246    fn decode_error_includes_call_context() {
247        let context = CallContext {
248            operation: "query_call",
249            canister_id: Principal::anonymous(),
250            caller: Principal::management_canister(),
251            method: "get",
252        };
253
254        let err = call_candid(
255            context,
256            (),
257            |_| Ok(vec![0xde, 0xad]),
258            |bytes| candid::decode_one::<u64>(bytes),
259        )
260        .expect_err("decode fails");
261
262        assert!(err.to_string().contains("candid decode failed"));
263        assert!(err.to_string().contains("operation=query_call"));
264        assert!(err.to_string().contains("method=get"));
265        assert!(err.to_string().contains("bytes=2"));
266        let CandidCallError::Decode { context, bytes, .. } = err else {
267            panic!("expected decode failure");
268        };
269        assert_eq!(bytes, 2);
270        assert_eq!(context.method, "get");
271    }
272
273    #[test]
274    fn reply_arguments_preserve_top_level_arity_and_wide_values() {
275        let context = CallContext {
276            operation: "update_call",
277            canister_id: Principal::anonymous(),
278            caller: Principal::management_canister(),
279            method: "measure",
280        };
281        let values = (u128::MAX, u128::from(u64::MAX) + 1);
282        let bytes = candid::encode_args(values).unwrap();
283        assert_eq!(candid::decode_args::<(u128, u128)>(&bytes).unwrap(), values);
284        assert!(candid::decode_one::<(u128, u128)>(&bytes).is_err());
285        let tuple = candid::encode_one(values).unwrap();
286        assert!(candid::decode_args::<(u128, u128)>(&tuple).is_err());
287        assert_eq!(candid::decode_one::<(u128, u128)>(&tuple).unwrap(), values);
288        let no_replies = candid::encode_args(()).unwrap();
289        candid::decode_args::<()>(&no_replies).unwrap();
290        let one_reply = candid::encode_args((7_u64,)).unwrap();
291        assert_eq!(candid::decode_args::<(u64,)>(&one_reply).unwrap(), (7,));
292        let result: (u64, bool, String) = call_candid(
293            context,
294            (42_u64, false),
295            |args| {
296                assert_eq!(
297                    candid::decode_args::<(u64, bool)>(&args).unwrap(),
298                    (42, false)
299                );
300                Ok(candid::encode_args((9_u64, true, "ready")).unwrap())
301            },
302            |bytes| candid::decode_args(bytes),
303        )
304        .unwrap();
305        assert_eq!(result, (9, true, "ready".into()));
306    }
307
308    #[test]
309    fn invalid_reply_arguments_keep_decode_context() {
310        let context = CallContext {
311            operation: "query_call",
312            canister_id: Principal::anonymous(),
313            caller: Principal::management_canister(),
314            method: "measure",
315        };
316        for bytes in [vec![0xde, 0xad], candid::encode_args((7_u64,)).unwrap()] {
317            let length = bytes.len();
318            let err = call_candid(
319                context,
320                (),
321                |_| Ok(bytes),
322                |bytes| candid::decode_args::<(u64, bool)>(bytes),
323            )
324            .unwrap_err();
325            assert!(err.to_string().contains(&format!("bytes={length}")));
326            let CandidCallError::Decode {
327                context: observed,
328                bytes,
329                ..
330            } = err
331            else {
332                panic!("expected decode failure");
333            };
334            assert_eq!(bytes, length);
335            assert_eq!(observed.operation, "query_call");
336            assert_eq!(observed.canister_id, context.canister_id);
337            assert_eq!(observed.caller, context.caller);
338            assert_eq!(observed.method, "measure");
339        }
340    }
341
342    #[test]
343    fn dead_instance_panic_is_classified_as_transport() {
344        let context = CallContext {
345            operation: "query_call",
346            canister_id: Principal::anonymous(),
347            caller: Principal::management_canister(),
348            method: "get",
349        };
350        let error = run_raw_call(context, || -> Result<Vec<u8>, pocket_ic::RejectResponse> {
351            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\" }}) }}");
352        })
353        .unwrap_err();
354
355        assert!(matches!(error, CandidCallError::Transport { .. }));
356    }
357
358    #[test]
359    fn unrelated_panics_resume_with_the_original_payload() {
360        for message in [
361            "application worker channel closed",
362            "fixture expected ConnectionRefused but observed another result",
363        ] {
364            let context = CallContext {
365                operation: "query_call",
366                canister_id: Principal::anonymous(),
367                caller: Principal::management_canister(),
368                method: "get",
369            };
370            let payload = message.to_owned();
371            let original_allocation = payload.as_ptr() as usize;
372            let result = std::panic::catch_unwind(|| {
373                run_raw_call(context, || -> Result<Vec<u8>, pocket_ic::RejectResponse> {
374                    std::panic::panic_any(payload);
375                })
376            });
377            let resumed = result
378                .expect_err("application panic must resume")
379                .downcast::<String>()
380                .expect("original panic type must be preserved");
381            assert_eq!(resumed.as_str(), message);
382            assert_eq!(resumed.as_ptr() as usize, original_allocation);
383        }
384    }
385}