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, transport};
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        self.update_candid_as(canister_id, Principal::anonymous(), method, args)
36    }
37
38    /// Execute [`update_candid`](Self::update_candid), panicking on harness errors.
39    ///
40    /// Application-level `Result<T, E>` values are returned without unwrapping.
41    #[track_caller]
42    fn update_candid_or_panic<T, A>(&self, canister_id: Principal, method: &str, args: A) -> T
43    where
44        T: CandidType + DeserializeOwned,
45        A: ArgumentEncoder,
46    {
47        self.update_candid(canister_id, method, args)
48            .unwrap_or_else(|err| panic!("{err}"))
49    }
50
51    /// Encode and execute an update call as an explicit caller.
52    fn update_candid_as<T, A>(
53        &self,
54        canister_id: Principal,
55        caller: Principal,
56        method: &str,
57        args: A,
58    ) -> Result<T, CandidCallError>
59    where
60        T: CandidType + DeserializeOwned,
61        A: ArgumentEncoder;
62
63    /// Execute [`update_candid_as`](Self::update_candid_as), panicking on harness errors.
64    ///
65    /// Application-level `Result<T, E>` values are returned without unwrapping.
66    #[track_caller]
67    fn update_candid_as_or_panic<T, A>(
68        &self,
69        canister_id: Principal,
70        caller: Principal,
71        method: &str,
72        args: A,
73    ) -> T
74    where
75        T: CandidType + DeserializeOwned,
76        A: ArgumentEncoder,
77    {
78        self.update_candid_as(canister_id, caller, method, args)
79            .unwrap_or_else(|err| panic!("{err}"))
80    }
81
82    /// Encode and execute an anonymous query call, then decode its result.
83    fn query_candid<T, A>(
84        &self,
85        canister_id: Principal,
86        method: &str,
87        args: A,
88    ) -> Result<T, CandidCallError>
89    where
90        T: CandidType + DeserializeOwned,
91        A: ArgumentEncoder,
92    {
93        self.query_candid_as(canister_id, Principal::anonymous(), method, args)
94    }
95
96    /// Execute [`query_candid`](Self::query_candid), panicking on harness errors.
97    ///
98    /// Application-level `Result<T, E>` values are returned without unwrapping.
99    #[track_caller]
100    fn query_candid_or_panic<T, A>(&self, canister_id: Principal, method: &str, args: A) -> T
101    where
102        T: CandidType + DeserializeOwned,
103        A: ArgumentEncoder,
104    {
105        self.query_candid(canister_id, method, args)
106            .unwrap_or_else(|err| panic!("{err}"))
107    }
108
109    /// Encode and execute a query call as an explicit caller.
110    fn query_candid_as<T, A>(
111        &self,
112        canister_id: Principal,
113        caller: Principal,
114        method: &str,
115        args: A,
116    ) -> Result<T, CandidCallError>
117    where
118        T: CandidType + DeserializeOwned,
119        A: ArgumentEncoder;
120
121    /// Execute [`query_candid_as`](Self::query_candid_as), panicking on harness errors.
122    ///
123    /// Application-level `Result<T, E>` values are returned without unwrapping.
124    #[track_caller]
125    fn query_candid_as_or_panic<T, A>(
126        &self,
127        canister_id: Principal,
128        caller: Principal,
129        method: &str,
130        args: A,
131    ) -> T
132    where
133        T: CandidType + DeserializeOwned,
134        A: ArgumentEncoder,
135    {
136        self.query_candid_as(canister_id, caller, method, args)
137            .unwrap_or_else(|err| panic!("{err}"))
138    }
139}
140
141impl CandidCallExt for PocketIc {
142    /// Generic update call helper with an explicit caller principal.
143    fn update_candid_as<T, A>(
144        &self,
145        canister_id: Principal,
146        caller: Principal,
147        method: &str,
148        args: A,
149    ) -> Result<T, CandidCallError>
150    where
151        T: CandidType + DeserializeOwned,
152        A: ArgumentEncoder,
153    {
154        let context = CallContext {
155            operation: "update_call",
156            canister_id,
157            caller,
158            method,
159        };
160        let bytes = encode_call_args(args, context)?;
161        let result = run_raw_call(context, || {
162            Self::update_call(self, canister_id, caller, method, bytes)
163        })?;
164
165        decode_call_result(&result, context)
166    }
167
168    /// Generic query call helper with an explicit caller principal.
169    fn query_candid_as<T, A>(
170        &self,
171        canister_id: Principal,
172        caller: Principal,
173        method: &str,
174        args: A,
175    ) -> Result<T, CandidCallError>
176    where
177        T: CandidType + DeserializeOwned,
178        A: ArgumentEncoder,
179    {
180        let context = CallContext {
181            operation: "query_call",
182            canister_id,
183            caller,
184            method,
185        };
186        let bytes = encode_call_args(args, context)?;
187        let result = run_raw_call(context, || {
188            Self::query_call(self, canister_id, caller, method, bytes)
189        })?;
190
191        decode_call_result(&result, context)
192    }
193}
194
195fn encode_call_args<A>(args: A, context: CallContext<'_>) -> Result<Vec<u8>, CandidCallError>
196where
197    A: ArgumentEncoder,
198{
199    encode_args(args).map_err(|err| CandidCallError::encode(context.to_error_context(), err))
200}
201
202fn decode_call_result<T>(result: &[u8], context: CallContext<'_>) -> Result<T, CandidCallError>
203where
204    T: CandidType + DeserializeOwned,
205{
206    decode_one(result)
207        .map_err(|err| CandidCallError::decode(context.to_error_context(), result.len(), err))
208}
209
210fn run_raw_call<F>(context: CallContext<'_>, call: F) -> Result<Vec<u8>, CandidCallError>
211where
212    F: FnOnce() -> Result<Vec<u8>, RejectResponse>,
213{
214    match catch_unwind(AssertUnwindSafe(call)) {
215        Ok(Ok(bytes)) => Ok(bytes),
216        Ok(Err(response)) => Err(CandidCallError::canister_reject(
217            context.to_error_context(),
218            response,
219        )),
220        Err(payload) if transport::panic_is_dead_instance_transport(payload.as_ref()) => {
221            Err(CandidCallError::transport(
222                context.to_error_context(),
223                transport::panic_payload_to_string(payload.as_ref()),
224            ))
225        }
226        Err(payload) => resume_unwind(payload),
227    }
228}
229
230#[cfg(test)]
231mod tests {
232    use candid::Principal;
233
234    use crate::pic::CandidCallErrorKind;
235
236    use super::{CallContext, decode_call_result, run_raw_call};
237
238    #[test]
239    fn decode_error_includes_call_context() {
240        let context = CallContext {
241            operation: "query_call",
242            canister_id: Principal::anonymous(),
243            caller: Principal::management_canister(),
244            method: "get",
245        };
246
247        let err = decode_call_result::<u64>(&[0xde, 0xad], context).expect_err("decode fails");
248
249        assert!(err.message().contains("candid decode_one failed"));
250        assert!(err.message().contains("operation=query_call"));
251        assert!(err.message().contains("method=get"));
252        assert!(err.message().contains("bytes=2"));
253        assert_eq!(err.kind(), CandidCallErrorKind::Decode);
254        assert_eq!(err.context().expect("decode error context").method(), "get");
255    }
256
257    #[test]
258    fn dead_instance_panic_is_classified_as_transport() {
259        let context = CallContext {
260            operation: "query_call",
261            canister_id: Principal::anonymous(),
262            caller: Principal::management_canister(),
263            method: "get",
264        };
265        let error = run_raw_call(context, || -> Result<Vec<u8>, pocket_ic::RejectResponse> {
266            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\" }}) }}");
267        })
268        .unwrap_err();
269
270        assert_eq!(error.kind(), crate::pic::CandidCallErrorKind::Transport);
271        assert!(error.reject_response().is_none());
272    }
273
274    #[test]
275    fn unrelated_panics_resume_with_the_original_payload() {
276        for message in [
277            "application worker channel closed",
278            "fixture expected ConnectionRefused but observed another result",
279        ] {
280            let context = CallContext {
281                operation: "query_call",
282                canister_id: Principal::anonymous(),
283                caller: Principal::management_canister(),
284                method: "get",
285            };
286            let payload = message.to_owned();
287            let original_allocation = payload.as_ptr() as usize;
288            let result = std::panic::catch_unwind(|| {
289                run_raw_call(context, || -> Result<Vec<u8>, pocket_ic::RejectResponse> {
290                    std::panic::panic_any(payload);
291                })
292            });
293            let resumed = result
294                .expect_err("application panic must resume")
295                .downcast::<String>()
296                .expect("original panic type must be preserved");
297            assert_eq!(resumed.as_str(), message);
298            assert_eq!(resumed.as_ptr() as usize, original_allocation);
299        }
300    }
301}