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