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, 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        if let Some(label) = &spec.label {
233            eprintln!("install_canister trapped for {canister_id} ({label})");
234        } else {
235            eprintln!("install_canister trapped for {canister_id}");
236        }
237        if let Ok(status) = pocket_ic.canister_status(canister_id, None) {
238            eprintln!("canister_status for {canister_id}: {status:?}");
239        }
240        if let Ok(logs) = pocket_ic.fetch_canister_logs(canister_id, Principal::anonymous()) {
241            for record in logs {
242                eprintln!("canister_log {canister_id}: {record:?}");
243            }
244        }
245        let message = startup::panic_payload_to_string(payload.as_ref());
246        return if let Some(label) = spec.label {
247            Err(CanisterInstallError::labeled(canister_id, label, message))
248        } else {
249            Err(CanisterInstallError::new(canister_id, message))
250        };
251    }
252
253    Ok(canister_id)
254}
255
256fn is_install_code_rate_limited(message: &str) -> bool {
257    message.contains("CanisterInstallCodeRateLimited")
258}
259
260fn retry_install_code_with<T, F, W>(
261    policy: RetryPolicy,
262    mut op: F,
263    mut wait_out_cooldown: W,
264) -> Result<T, String>
265where
266    F: FnMut() -> Result<T, String>,
267    W: FnMut(),
268{
269    for attempt in 1..=policy.max_attempts() {
270        match op() {
271            Ok(value) => return Ok(value),
272            Err(err) if is_install_code_rate_limited(&err) && attempt < policy.max_attempts() => {
273                wait_out_cooldown();
274            }
275            Err(err) => return Err(err),
276        }
277    }
278
279    unreachable!("RetryPolicy guarantees at least one attempt")
280}
281
282#[cfg(test)]
283mod tests {
284    use std::{cell::Cell, time::Duration};
285
286    use super::{RetryPolicy, retry_install_code_with};
287
288    const RATE_LIMITED: &str = "CanisterInstallCodeRateLimited";
289
290    #[test]
291    fn retry_policy_counts_the_first_attempt() {
292        let attempts = Cell::new(0);
293        let waits = Cell::new(0);
294        let result = retry_install_code_with(
295            RetryPolicy::new(3, Duration::from_secs(1)),
296            || {
297                attempts.set(attempts.get() + 1);
298                Err::<(), _>(RATE_LIMITED.to_string())
299            },
300            || waits.set(waits.get() + 1),
301        );
302
303        assert_eq!(result, Err(RATE_LIMITED.to_string()));
304        assert_eq!(attempts.get(), 3);
305        assert_eq!(waits.get(), 2);
306    }
307
308    #[test]
309    fn retry_policy_stops_on_non_rate_limit_failure() {
310        let attempts = Cell::new(0);
311        let result = retry_install_code_with(
312            RetryPolicy::new(3, Duration::from_secs(1)),
313            || {
314                attempts.set(attempts.get() + 1);
315                Err::<(), _>("not retryable".to_string())
316            },
317            || panic!("non-rate-limit failure must not wait"),
318        );
319
320        assert_eq!(result, Err("not retryable".to_string()));
321        assert_eq!(attempts.get(), 1);
322    }
323}
ic_testkit/pic/lifecycle.rs

ic_testkit/pic/
lifecycle.rs
