ic_testkit/pic/
lifecycle.rs

1use std::panic::{AssertUnwindSafe, catch_unwind};
2use std::time::Duration;
3
4use candid::Principal;
5use pocket_ic::PocketIc;
6
7use super::{CanisterInstallError, PocketIcDiagnosticsExt, PocketIcTimeExt, startup};
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, String>
129    where
130        F: FnMut() -> Result<T, String>;
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_n(2);
206    }
207
208    fn retry_install_code<T, F>(&self, policy: RetryPolicy, op: F) -> Result<T, String>
209    where
210        F: FnMut() -> Result<T, String>,
211    {
212        retry_install_code_with(policy, op, || {
213            self.wait_out_install_code_rate_limit(policy.cooldown());
214        })
215    }
216}
217
218// Install a canister after creating it and optionally adding extra cycles.
219fn try_create_funded_and_install(
220    pocket_ic: &PocketIc,
221    spec: InstallSpec,
222) -> Result<Principal, CanisterInstallError> {
223    let canister_id = pocket_ic.create_canister();
224    if spec.cycles > 0 {
225        let _ = pocket_ic.add_cycles(canister_id, spec.cycles);
226    }
227
228    let install = catch_unwind(AssertUnwindSafe(|| {
229        pocket_ic.install_canister(canister_id, spec.wasm, spec.init_bytes, spec.install_sender);
230    }));
231    if let Err(payload) = install {
232        let message = startup::panic_payload_to_string(payload.as_ref());
233        let context = if let Some(label) = &spec.label {
234            format!("install_canister trapped ({label})")
235        } else {
236            "install_canister trapped".to_string()
237        };
238        // Diagnostics are best-effort and must never replace the original
239        // structured install failure, including if stderr or PocketIC fails.
240        let _ = catch_unwind(AssertUnwindSafe(|| {
241            pocket_ic.dump_canister_debug(canister_id, &context);
242        }));
243
244        return if let Some(label) = spec.label {
245            Err(CanisterInstallError::labeled(canister_id, label, message))
246        } else {
247            Err(CanisterInstallError::new(canister_id, message))
248        };
249    }
250
251    Ok(canister_id)
252}
253
254fn is_install_code_rate_limited(message: &str) -> bool {
255    message.contains("CanisterInstallCodeRateLimited")
256}
257
258fn retry_install_code_with<T, F, W>(
259    policy: RetryPolicy,
260    mut op: F,
261    mut wait_out_cooldown: W,
262) -> Result<T, String>
263where
264    F: FnMut() -> Result<T, String>,
265    W: FnMut(),
266{
267    for attempt in 1..=policy.max_attempts() {
268        match op() {
269            Ok(value) => return Ok(value),
270            Err(err) if is_install_code_rate_limited(&err) && attempt < policy.max_attempts() => {
271                wait_out_cooldown();
272            }
273            Err(err) => return Err(err),
274        }
275    }
276
277    unreachable!("RetryPolicy guarantees at least one attempt")
278}
279
280#[cfg(test)]
281mod tests {
282    use std::{cell::Cell, time::Duration};
283
284    use super::{RetryPolicy, retry_install_code_with};
285
286    const RATE_LIMITED: &str = "CanisterInstallCodeRateLimited";
287
288    #[test]
289    fn retry_policy_counts_the_first_attempt() {
290        let attempts = Cell::new(0);
291        let waits = Cell::new(0);
292        let result = retry_install_code_with(
293            RetryPolicy::new(3, Duration::from_secs(1)),
294            || {
295                attempts.set(attempts.get() + 1);
296                Err::<(), _>(RATE_LIMITED.to_string())
297            },
298            || waits.set(waits.get() + 1),
299        );
300
301        assert_eq!(result, Err(RATE_LIMITED.to_string()));
302        assert_eq!(attempts.get(), 3);
303        assert_eq!(waits.get(), 2);
304    }
305
306    #[test]
307    fn retry_policy_stops_on_non_rate_limit_failure() {
308        let attempts = Cell::new(0);
309        let result = retry_install_code_with(
310            RetryPolicy::new(3, Duration::from_secs(1)),
311            || {
312                attempts.set(attempts.get() + 1);
313                Err::<(), _>("not retryable".to_string())
314            },
315            || panic!("non-rate-limit failure must not wait"),
316        );
317
318        assert_eq!(result, Err("not retryable".to_string()));
319        assert_eq!(attempts.get(), 1);
320    }
321}
ic_testkit/pic/lifecycle.rs

ic_testkit/pic/
lifecycle.rs
