Skip to main content

ic_testkit/pic/
calls.rs

1use candid::{CandidType, Principal, decode_one, encode_args, utils::ArgumentEncoder};
2use serde::de::DeserializeOwned;
3
4use super::{Pic, PicCallContext, PicCallError};
5
6#[derive(Clone, Copy)]
7struct CallContext<'a> {
8    operation: &'static str,
9    canister_id: Principal,
10    caller: Principal,
11    method: &'a str,
12}
13
14impl CallContext<'_> {
15    fn to_error_context(self) -> PicCallContext {
16        PicCallContext::new(self.operation, self.canister_id, self.caller, self.method)
17    }
18}
19
20impl Pic {
21    /// Generic update call helper (serializes args + decodes result).
22    pub fn update_call<T, A>(
23        &self,
24        canister_id: Principal,
25        method: &str,
26        args: A,
27    ) -> Result<T, PicCallError>
28    where
29        T: CandidType + DeserializeOwned,
30        A: ArgumentEncoder,
31    {
32        self.update_call_as(canister_id, Principal::anonymous(), method, args)
33    }
34
35    /// Generic update call helper that panics on transport or Candid codec failure.
36    ///
37    /// This does not unwrap application-level results. For example,
38    /// `update_call_or_panic::<Result<T, E>, _>(...)` returns `Result<T, E>`.
39    #[track_caller]
40    pub fn update_call_or_panic<T, A>(&self, canister_id: Principal, method: &str, args: A) -> T
41    where
42        T: CandidType + DeserializeOwned,
43        A: ArgumentEncoder,
44    {
45        self.update_call(canister_id, method, args)
46            .unwrap_or_else(|err| panic!("{err}"))
47    }
48
49    /// Generic update call helper with an explicit caller principal.
50    pub fn update_call_as<T, A>(
51        &self,
52        canister_id: Principal,
53        caller: Principal,
54        method: &str,
55        args: A,
56    ) -> Result<T, PicCallError>
57    where
58        T: CandidType + DeserializeOwned,
59        A: ArgumentEncoder,
60    {
61        let context = CallContext {
62            operation: "update_call",
63            canister_id,
64            caller,
65            method,
66        };
67        let bytes = encode_call_args(args, context)?;
68        let result = self
69            .inner
70            .update_call(canister_id, caller, method, bytes)
71            .map_err(|err| PicCallError::transport(context.to_error_context(), err))?;
72
73        decode_call_result(&result, context)
74    }
75
76    /// Generic update call helper with an explicit caller principal that panics
77    /// on transport or Candid codec failure.
78    ///
79    /// This does not unwrap application-level results. For example,
80    /// `update_call_as_or_panic::<Result<T, E>, _>(...)` returns `Result<T, E>`.
81    #[track_caller]
82    pub fn update_call_as_or_panic<T, A>(
83        &self,
84        canister_id: Principal,
85        caller: Principal,
86        method: &str,
87        args: A,
88    ) -> T
89    where
90        T: CandidType + DeserializeOwned,
91        A: ArgumentEncoder,
92    {
93        self.update_call_as(canister_id, caller, method, args)
94            .unwrap_or_else(|err| panic!("{err}"))
95    }
96
97    /// Generic query call helper.
98    pub fn query_call<T, A>(
99        &self,
100        canister_id: Principal,
101        method: &str,
102        args: A,
103    ) -> Result<T, PicCallError>
104    where
105        T: CandidType + DeserializeOwned,
106        A: ArgumentEncoder,
107    {
108        self.query_call_as(canister_id, Principal::anonymous(), method, args)
109    }
110
111    /// Generic query call helper that panics on transport or Candid codec failure.
112    ///
113    /// This does not unwrap application-level results. For example,
114    /// `query_call_or_panic::<Result<T, E>, _>(...)` returns `Result<T, E>`.
115    #[track_caller]
116    pub fn query_call_or_panic<T, A>(&self, canister_id: Principal, method: &str, args: A) -> T
117    where
118        T: CandidType + DeserializeOwned,
119        A: ArgumentEncoder,
120    {
121        self.query_call(canister_id, method, args)
122            .unwrap_or_else(|err| panic!("{err}"))
123    }
124
125    /// Generic query call helper with an explicit caller principal.
126    pub fn query_call_as<T, A>(
127        &self,
128        canister_id: Principal,
129        caller: Principal,
130        method: &str,
131        args: A,
132    ) -> Result<T, PicCallError>
133    where
134        T: CandidType + DeserializeOwned,
135        A: ArgumentEncoder,
136    {
137        let context = CallContext {
138            operation: "query_call",
139            canister_id,
140            caller,
141            method,
142        };
143        let bytes = encode_call_args(args, context)?;
144        let result = self
145            .inner
146            .query_call(canister_id, caller, method, bytes)
147            .map_err(|err| PicCallError::transport(context.to_error_context(), err))?;
148
149        decode_call_result(&result, context)
150    }
151
152    /// Generic query call helper with an explicit caller principal that panics
153    /// on transport or Candid codec failure.
154    ///
155    /// This does not unwrap application-level results. For example,
156    /// `query_call_as_or_panic::<Result<T, E>, _>(...)` returns `Result<T, E>`.
157    #[track_caller]
158    pub fn query_call_as_or_panic<T, A>(
159        &self,
160        canister_id: Principal,
161        caller: Principal,
162        method: &str,
163        args: A,
164    ) -> T
165    where
166        T: CandidType + DeserializeOwned,
167        A: ArgumentEncoder,
168    {
169        self.query_call_as(canister_id, caller, method, args)
170            .unwrap_or_else(|err| panic!("{err}"))
171    }
172
173    /// Advance PocketIC by a fixed number of ticks.
174    pub fn tick_n(&self, times: usize) {
175        for _ in 0..times {
176            self.tick();
177        }
178    }
179}
180
181fn encode_call_args<A>(args: A, context: CallContext<'_>) -> Result<Vec<u8>, PicCallError>
182where
183    A: ArgumentEncoder,
184{
185    encode_args(args).map_err(|err| PicCallError::encode(context.to_error_context(), err))
186}
187
188fn decode_call_result<T>(result: &[u8], context: CallContext<'_>) -> Result<T, PicCallError>
189where
190    T: CandidType + DeserializeOwned,
191{
192    decode_one(result)
193        .map_err(|err| PicCallError::decode(context.to_error_context(), result.len(), err))
194}
195
196#[cfg(test)]
197mod tests {
198    use candid::Principal;
199
200    use crate::pic::PicCallErrorKind;
201
202    use super::{CallContext, decode_call_result};
203
204    #[test]
205    fn decode_error_includes_call_context() {
206        let context = CallContext {
207            operation: "query_call",
208            canister_id: Principal::anonymous(),
209            caller: Principal::management_canister(),
210            method: "get",
211        };
212
213        let err = decode_call_result::<u64>(&[0xde, 0xad], context).expect_err("decode fails");
214
215        assert!(err.message().contains("candid decode_one failed"));
216        assert!(err.message().contains("operation=query_call"));
217        assert!(err.message().contains("method=get"));
218        assert!(err.message().contains("bytes=2"));
219        assert_eq!(err.kind(), PicCallErrorKind::Decode);
220        assert_eq!(err.context().expect("decode error context").method(), "get");
221    }
222}