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.
35#[expect(
36    async_fn_in_trait,
37    reason = "extension futures may retain caller-owned non-Send callbacks, iterators or writers; pool recipes separately require Send"
38)]
39pub trait CanisterInstallExt {
40    /// Create and install one canister from a reusable specification.
41    ///
42    /// Retains typed client failures during creation, cycle funding, and installation.
43    /// The error identifies the stage and retains an id when creation succeeded.
44    /// Callers choose whether to collect diagnostics; this performs no diagnostic
45    /// requests or error logging.
46    ///
47    /// # Errors
48    /// Returns the failed creation or installation stage, retaining canister identity when available.
49    async fn try_create_and_install(
50        &self,
51        spec: InstallSpec,
52    ) -> Result<Principal, CanisterInstallError>;
53
54    /// Advance simulated time once, then tick twice for an install-code cooldown.
55    /// Zero duration skips time reads/writes while retaining both ticks.
56    ///
57    /// # Errors
58    /// Returns the first client failure without dispatching remaining ticks.
59    async fn wait_out_install_code_rate_limit(&self, cooldown: Duration)
60    -> Result<(), ClientError>;
61
62    /// Retry only while `PocketIC` reports install-code rate limiting.
63    ///
64    /// `RetryPolicy::max_attempts` includes the initial call. Before each
65    /// retry, this advances simulated time by the configured cooldown and
66    /// executes two `PocketIC` ticks. Other rejections are returned unchanged.
67    ///
68    /// # Errors
69    /// Returns non-rate-limit failures immediately, or the final rate-limit rejection.
70    /// A cooldown failure stops further installation attempts.
71    async fn retry_install_code<T, F>(&self, policy: RetryPolicy, op: F) -> Result<T, ClientError>
72    where
73        F: AsyncFnMut() -> Result<T, ClientError>;
74}
75
76impl CanisterInstallExt for PocketIc {
77    /// Install one arbitrary wasm module from a generic install specification.
78    async fn try_create_and_install(
79        &self,
80        spec: InstallSpec,
81    ) -> Result<Principal, CanisterInstallError> {
82        let canister_id = self.create_canister().await.map_err(|source| {
83            CanisterInstallError::CreateCanister {
84                label: spec.label.clone(),
85                source,
86            }
87        })?;
88        if spec.cycles > 0 {
89            self.add_cycles(canister_id, spec.cycles)
90                .await
91                .map_err(|source| CanisterInstallError::AddCycles {
92                    canister_id,
93                    label: spec.label.clone(),
94                    source,
95                })?;
96        }
97
98        // Move the Wasm and init payload while retaining the label for failures.
99        self.install_canister(canister_id, spec.wasm, spec.init_bytes, spec.install_sender)
100            .await
101            .map_err(|source| CanisterInstallError::InstallCode {
102                canister_id,
103                label: spec.label,
104                source,
105            })?;
106
107        Ok(canister_id)
108    }
109
110    /// Wait out the `PocketIC` `install_code` cooldown window inside the same instance.
111    async fn wait_out_install_code_rate_limit(
112        &self,
113        cooldown: Duration,
114    ) -> Result<(), ClientError> {
115        super::time::advance_and_tick(self, cooldown, 2).await
116    }
117
118    async fn retry_install_code<T, F>(&self, policy: RetryPolicy, op: F) -> Result<T, ClientError>
119    where
120        F: AsyncFnMut() -> Result<T, ClientError>,
121    {
122        retry_install_code_with(policy, op, async || {
123            self.wait_out_install_code_rate_limit(policy.cooldown).await
124        })
125        .await
126    }
127}
128
129async fn retry_install_code_with<T, F, W>(
130    policy: RetryPolicy,
131    mut op: F,
132    mut wait_out_cooldown: W,
133) -> Result<T, ClientError>
134where
135    F: AsyncFnMut() -> Result<T, ClientError>,
136    W: AsyncFnMut() -> Result<(), ClientError>,
137{
138    let mut retries = 1..policy.max_attempts.get();
139    loop {
140        match op().await {
141            Err(ClientError::Canister(err))
142                if err.error_code == ErrorCode::CanisterInstallCodeRateLimited
143                    && retries.next().is_some() =>
144            {
145                wait_out_cooldown().await?;
146            }
147            result => return result,
148        }
149    }
150}
151
152#[cfg(test)]
153mod tests;
ic_testkit/pic/lifecycle.rs

ic_testkit/pic/
lifecycle.rs
