ic_testkit/pic/
lifecycle.rs

1use std::panic::{AssertUnwindSafe, catch_unwind};
2use std::time::Duration;
3
4use candid::Principal;
5use pocket_ic::{ErrorCode, PocketIc, RejectResponse};
6
7use super::{CanisterInstallError, PocketIcDiagnosticsExt, transport};
8
9///
10/// InstallSpec
11///
12
13#[non_exhaustive]
14pub struct InstallSpec {
15    pub wasm: Vec<u8>,
16    pub init_bytes: Vec<u8>,
17    pub cycles: u128,
18    pub install_sender: Option<Principal>,
19    pub label: Option<String>,
20}
21
22impl InstallSpec {
23    /// Build one generic canister install specification.
24    #[must_use]
25    pub const fn new(wasm: Vec<u8>, init_bytes: Vec<u8>, cycles: u128) -> Self {
26        Self {
27            wasm,
28            init_bytes,
29            cycles,
30            install_sender: None,
31            label: None,
32        }
33    }
34
35    /// Set the management-call sender used for `install_canister`.
36    #[must_use]
37    pub const fn install_sender(mut self, sender: Principal) -> Self {
38        self.install_sender = Some(sender);
39        self
40    }
41
42    /// Set a diagnostic label for install failures.
43    #[must_use]
44    pub fn label(mut self, label: impl Into<String>) -> Self {
45        self.label = Some(label.into());
46        self
47    }
48}
49
50/// Retry limits and simulated cooldown for install-code operations.
51#[derive(Clone, Copy, Debug, Eq, PartialEq)]
52pub struct RetryPolicy {
53    max_attempts: usize,
54    cooldown: Duration,
55}
56
57impl RetryPolicy {
58    /// Create a policy with an exact, non-zero maximum attempt count.
59    #[must_use]
60    pub const fn new(max_attempts: usize, cooldown: Duration) -> Self {
61        assert!(
62            max_attempts > 0,
63            "retry policy requires at least one attempt"
64        );
65        Self {
66            max_attempts,
67            cooldown,
68        }
69    }
70
71    /// Read the maximum number of operation attempts, including the first.
72    #[must_use]
73    pub const fn max_attempts(self) -> usize {
74        self.max_attempts
75    }
76
77    /// Read the simulated cooldown applied between rate-limited attempts.
78    #[must_use]
79    pub const fn cooldown(self) -> Duration {
80        self.cooldown
81    }
82}
83
84/// Generic canister installation and install-code retry policies.
85pub trait CanisterInstallExt {
86    /// Create and install one canister from raw wasm and init bytes.
87    #[must_use]
88    fn create_and_install_with_args(
89        &self,
90        wasm: Vec<u8>,
91        init_bytes: Vec<u8>,
92        install_cycles: u128,
93    ) -> Principal;
94
95    /// Fallible counterpart to [`create_and_install_with_args`](Self::create_and_install_with_args).
96    fn try_create_and_install_with_args(
97        &self,
98        wasm: Vec<u8>,
99        init_bytes: Vec<u8>,
100        install_cycles: u128,
101    ) -> Result<Principal, CanisterInstallError>;
102
103    /// Create and install one canister from a reusable specification.
104    #[must_use]
105    fn create_and_install(&self, spec: InstallSpec) -> Principal;
106
107    /// Fallible counterpart to [`create_and_install`](Self::create_and_install).
108    fn try_create_and_install(&self, spec: InstallSpec) -> Result<Principal, CanisterInstallError>;
109
110    /// Sequentially create and install multiple canisters.
111    #[must_use]
112    fn create_and_install_many<I>(&self, specs: I) -> Vec<Principal>
113    where
114        I: IntoIterator<Item = InstallSpec>;
115
116    /// Fallible counterpart to [`create_and_install_many`](Self::create_and_install_many).
117    fn try_create_and_install_many<I>(
118        &self,
119        specs: I,
120    ) -> Result<Vec<Principal>, CanisterInstallError>
121    where
122        I: IntoIterator<Item = InstallSpec>;
123
124    /// Advance simulated time and rounds past an install-code cooldown.
125    fn wait_out_install_code_rate_limit(&self, cooldown: Duration);
126
127    /// Retry an operation only while PocketIC reports install-code rate limiting.
128    fn retry_install_code<T, F>(&self, policy: RetryPolicy, op: F) -> Result<T, RejectResponse>
129    where
130        F: FnMut() -> Result<T, RejectResponse>;
131}
132
133impl CanisterInstallExt for PocketIc {
134    /// Install one arbitrary wasm module with caller-provided init bytes.
135    ///
136    /// This is the generic install path for downstreams that use `ic-testkit`
137    /// without depending on application-specific init payload conventions.
138    fn create_and_install_with_args(
139        &self,
140        wasm: Vec<u8>,
141        init_bytes: Vec<u8>,
142        install_cycles: u128,
143    ) -> Principal {
144        self.try_create_and_install_with_args(wasm, init_bytes, install_cycles)
145            .unwrap_or_else(|err| panic!("{err}"))
146    }
147
148    /// Install one arbitrary wasm module with caller-provided init bytes.
149    fn try_create_and_install_with_args(
150        &self,
151        wasm: Vec<u8>,
152        init_bytes: Vec<u8>,
153        install_cycles: u128,
154    ) -> Result<Principal, CanisterInstallError> {
155        self.try_create_and_install(InstallSpec::new(wasm, init_bytes, install_cycles))
156    }
157
158    /// Install one arbitrary wasm module from a generic install specification.
159    fn create_and_install(&self, spec: InstallSpec) -> Principal {
160        self.try_create_and_install(spec)
161            .unwrap_or_else(|err| panic!("{err}"))
162    }
163
164    /// Install one arbitrary wasm module from a generic install specification.
165    fn try_create_and_install(&self, spec: InstallSpec) -> Result<Principal, CanisterInstallError> {
166        try_create_funded_and_install(self, spec)
167    }
168
169    /// Sequentially install multiple arbitrary wasm modules into this PocketIC instance.
170    ///
171    /// Installs are attempted in iterator order. If one install fails, earlier
172    /// installs remain in the PocketIC instance, the failed canister may exist
173    /// with the id exposed by `CanisterInstallError::canister_id()`, and later
174    /// installs are not attempted.
175    fn create_and_install_many<I>(&self, specs: I) -> Vec<Principal>
176    where
177        I: IntoIterator<Item = InstallSpec>,
178    {
179        self.try_create_and_install_many(specs)
180            .unwrap_or_else(|err| panic!("{err}"))
181    }
182
183    /// Sequentially install multiple arbitrary wasm modules into this PocketIC instance.
184    ///
185    /// Installs are attempted in iterator order. If one install fails, earlier
186    /// installs remain in the PocketIC instance, the failed canister may exist
187    /// with the id exposed by `CanisterInstallError::canister_id()`, and later
188    /// installs are not attempted.
189    fn try_create_and_install_many<I>(
190        &self,
191        specs: I,
192    ) -> Result<Vec<Principal>, CanisterInstallError>
193    where
194        I: IntoIterator<Item = InstallSpec>,
195    {
196        specs
197            .into_iter()
198            .map(|spec| self.try_create_and_install(spec))
199            .collect()
200    }
201
202    /// Wait out the PocketIC `install_code` cooldown window inside the same instance.
203    fn wait_out_install_code_rate_limit(&self, cooldown: Duration) {
204        self.advance_time(cooldown);
205        self.tick();
206        self.tick();
207    }
208
209    fn retry_install_code<T, F>(&self, policy: RetryPolicy, op: F) -> Result<T, RejectResponse>
210    where
211        F: FnMut() -> Result<T, RejectResponse>,
212    {
213        retry_install_code_with(policy, op, || {
214            self.wait_out_install_code_rate_limit(policy.cooldown());
215        })
216    }
217}
218
219// Install a canister after creating it and optionally adding extra cycles.
220fn try_create_funded_and_install(
221    pocket_ic: &PocketIc,
222    spec: InstallSpec,
223) -> Result<Principal, CanisterInstallError> {
224    let canister_id = pocket_ic.create_canister();
225    if spec.cycles > 0 {
226        let _ = pocket_ic.add_cycles(canister_id, spec.cycles);
227    }
228
229    let install = catch_unwind(AssertUnwindSafe(|| {
230        pocket_ic.install_canister(canister_id, spec.wasm, spec.init_bytes, spec.install_sender);
231    }));
232    if let Err(payload) = install {
233        let message = transport::panic_payload_to_string(payload.as_ref());
234        let context = if let Some(label) = &spec.label {
235            format!("install_canister trapped ({label})")
236        } else {
237            "install_canister trapped".to_string()
238        };
239        // Diagnostics are best-effort and must never replace the original
240        // structured install failure, including if stderr or PocketIC fails.
241        let _ = catch_unwind(AssertUnwindSafe(|| {
242            pocket_ic.dump_canister_debug(canister_id, &context);
243        }));
244
245        return if let Some(label) = spec.label {
246            Err(CanisterInstallError::labeled(canister_id, label, message))
247        } else {
248            Err(CanisterInstallError::new(canister_id, message))
249        };
250    }
251
252    Ok(canister_id)
253}
254
255fn is_install_code_rate_limited(response: &RejectResponse) -> bool {
256    response.error_code == ErrorCode::CanisterInstallCodeRateLimited
257}
258
259fn retry_install_code_with<T, F, W>(
260    policy: RetryPolicy,
261    mut op: F,
262    mut wait_out_cooldown: W,
263) -> Result<T, RejectResponse>
264where
265    F: FnMut() -> Result<T, RejectResponse>,
266    W: FnMut(),
267{
268    for attempt in 1..=policy.max_attempts() {
269        match op() {
270            Ok(value) => return Ok(value),
271            Err(err) if is_install_code_rate_limited(&err) && attempt < policy.max_attempts() => {
272                wait_out_cooldown();
273            }
274            Err(err) => return Err(err),
275        }
276    }
277
278    unreachable!("RetryPolicy guarantees at least one attempt")
279}
280
281#[cfg(test)]
282mod tests {
283    use std::{cell::Cell, time::Duration};
284
285    use pocket_ic::{ErrorCode, RejectCode, RejectResponse};
286
287    use super::{RetryPolicy, retry_install_code_with};
288
289    fn rejection(error_code: ErrorCode, message: &str) -> RejectResponse {
290        RejectResponse {
291            reject_code: RejectCode::SysTransient,
292            reject_message: message.to_string(),
293            error_code,
294            certified: false,
295        }
296    }
297
298    #[test]
299    fn retry_policy_counts_the_first_attempt() {
300        let attempts = Cell::new(0);
301        let waits = Cell::new(0);
302        let rate_limited = rejection(
303            ErrorCode::CanisterInstallCodeRateLimited,
304            "install-code rate limit",
305        );
306        let result = retry_install_code_with(
307            RetryPolicy::new(3, Duration::from_secs(1)),
308            || {
309                attempts.set(attempts.get() + 1);
310                Err::<(), _>(rate_limited.clone())
311            },
312            || waits.set(waits.get() + 1),
313        );
314
315        assert_eq!(result, Err(rate_limited));
316        assert_eq!(attempts.get(), 3);
317        assert_eq!(waits.get(), 2);
318    }
319
320    #[test]
321    fn retry_policy_stops_on_non_rate_limit_failure() {
322        let attempts = Cell::new(0);
323        let not_retryable = rejection(ErrorCode::CanisterRejectedMessage, "not retryable");
324        let result = retry_install_code_with(
325            RetryPolicy::new(3, Duration::from_secs(1)),
326            || {
327                attempts.set(attempts.get() + 1);
328                Err::<(), _>(not_retryable.clone())
329            },
330            || panic!("non-rate-limit failure must not wait"),
331        );
332
333        assert_eq!(result, Err(not_retryable));
334        assert_eq!(attempts.get(), 1);
335    }
336}
ic_testkit/pic/lifecycle.rs

ic_testkit/pic/
lifecycle.rs
