ic_testkit/pic/
lifecycle.rs

1use std::{num::NonZeroUsize, time::Duration};
2
3use crate::pocket_ic::{ClientError, ErrorCode, PocketIc};
4use ic_platform_types::principal::ExternalPrincipal as Principal;
5
6use super::CanisterInstallError;
7
8/// Inputs and diagnostic context for one generic canister installation.
9pub struct InstallSpec {
10    /// Wasm module installed into the newly created canister.
11    pub wasm: Vec<u8>,
12    /// Raw initialization bytes encoded by the caller.
13    pub init_bytes: Vec<u8>,
14    /// Extra cycles added after `PocketIC` creates the canister and before install.
15    pub cycles: u128,
16    /// Optional sender for the management-canister install operation.
17    pub install_sender: Option<Principal>,
18    /// Optional human-readable context included in install errors.
19    pub label: Option<String>,
20}
21
22/// Retry limits and simulated cooldown for install-code operations.
23#[derive(Clone, Copy, Debug, Eq, PartialEq)]
24pub struct RetryPolicy {
25    /// Maximum attempts, including the initial call.
26    pub max_attempts: NonZeroUsize,
27    /// Simulated cooldown between retryable attempts.
28    pub cooldown: Duration,
29}
30
31/// Generic canister installation and structured install-code retry policy.
32///
33/// The extension creates canisters using `PocketIC` defaults. `InstallSpec::cycles`
34/// is an additional top-up, not the complete initial balance.
35pub trait CanisterInstallExt {
36    /// Create and install one canister from a reusable specification.
37    ///
38    /// Retains typed client failures during creation, cycle funding, and installation.
39    /// The error identifies the stage and retains an id when creation succeeded.
40    /// Callers choose whether to collect diagnostics; this performs no diagnostic
41    /// requests or error logging.
42    ///
43    /// # Errors
44    /// Returns the failed creation or installation stage, retaining canister identity when available.
45    fn try_create_and_install(&self, spec: InstallSpec) -> Result<Principal, CanisterInstallError>;
46
47    /// Advance simulated time once, then tick twice for an install-code cooldown.
48    /// Zero duration skips time reads/writes while retaining both ticks.
49    ///
50    /// # Errors
51    /// Returns the first client failure without dispatching remaining ticks.
52    fn wait_out_install_code_rate_limit(&self, cooldown: Duration) -> Result<(), ClientError>;
53
54    /// Retry only while `PocketIC` reports install-code rate limiting.
55    ///
56    /// `RetryPolicy::max_attempts` includes the initial call. Before each
57    /// retry, this advances simulated time by the configured cooldown and
58    /// executes two `PocketIC` ticks. Other rejections are returned unchanged.
59    ///
60    /// # Errors
61    /// Returns non-rate-limit failures immediately, or the final rate-limit rejection.
62    /// A cooldown failure stops further installation attempts.
63    fn retry_install_code<T, F>(&self, policy: RetryPolicy, op: F) -> Result<T, ClientError>
64    where
65        F: FnMut() -> Result<T, ClientError>;
66}
67
68impl CanisterInstallExt for PocketIc {
69    /// Install one arbitrary wasm module from a generic install specification.
70    fn try_create_and_install(&self, spec: InstallSpec) -> Result<Principal, CanisterInstallError> {
71        let canister_id =
72            self.create_canister()
73                .map_err(|source| CanisterInstallError::CreateCanister {
74                    label: spec.label.clone(),
75                    source,
76                })?;
77        if spec.cycles > 0 {
78            self.add_cycles(canister_id, spec.cycles)
79                .map_err(|source| CanisterInstallError::AddCycles {
80                    canister_id,
81                    label: spec.label.clone(),
82                    source,
83                })?;
84        }
85
86        // Move the Wasm and init payload while retaining the label for failures.
87        self.install_canister(canister_id, spec.wasm, spec.init_bytes, spec.install_sender)
88            .map_err(|source| CanisterInstallError::InstallCode {
89                canister_id,
90                label: spec.label,
91                source,
92            })?;
93
94        Ok(canister_id)
95    }
96
97    /// Wait out the `PocketIC` `install_code` cooldown window inside the same instance.
98    fn wait_out_install_code_rate_limit(&self, cooldown: Duration) -> Result<(), ClientError> {
99        super::time::advance_and_tick(self, cooldown, 2)
100    }
101
102    fn retry_install_code<T, F>(&self, policy: RetryPolicy, op: F) -> Result<T, ClientError>
103    where
104        F: FnMut() -> Result<T, ClientError>,
105    {
106        retry_install_code_with(policy, op, || {
107            self.wait_out_install_code_rate_limit(policy.cooldown)
108        })
109    }
110}
111
112fn retry_install_code_with<T, F, W>(
113    policy: RetryPolicy,
114    mut op: F,
115    mut wait_out_cooldown: W,
116) -> Result<T, ClientError>
117where
118    F: FnMut() -> Result<T, ClientError>,
119    W: FnMut() -> Result<(), ClientError>,
120{
121    let mut retries = 1..policy.max_attempts.get();
122    loop {
123        match op() {
124            Err(ClientError::Canister(err))
125                if err.error_code == ErrorCode::CanisterInstallCodeRateLimited
126                    && retries.next().is_some() =>
127            {
128                wait_out_cooldown()?;
129            }
130            result => return result,
131        }
132    }
133}
134
135#[cfg(test)]
136mod tests;
ic_testkit/pic/lifecycle.rs

ic_testkit/pic/
lifecycle.rs
