ic_testkit/pic/
lifecycle.rs

1use std::{num::NonZeroUsize, time::Duration};
2
3use candid::Principal;
4use pocket_ic::{ErrorCode, PocketIc, RejectResponse};
5
6use super::{CanisterInstallError, transport::capture_operation};
7
8/// Inputs and diagnostic context for one generic canister installation.
9#[non_exhaustive]
10pub struct InstallSpec {
11    /// Wasm module installed into the newly created canister.
12    pub wasm: Vec<u8>,
13    /// Raw Candid or application-specific initialization bytes.
14    pub init_bytes: Vec<u8>,
15    /// Extra cycles added after PocketIC creates the canister and before install.
16    pub cycles: u128,
17    /// Optional sender for the management-canister install operation.
18    pub install_sender: Option<Principal>,
19    /// Optional human-readable context included in install errors.
20    pub label: Option<String>,
21}
22
23impl InstallSpec {
24    /// Build one generic canister install specification.
25    #[must_use]
26    pub const fn new(wasm: Vec<u8>, init_bytes: Vec<u8>, cycles: u128) -> Self {
27        Self {
28            wasm,
29            init_bytes,
30            cycles,
31            install_sender: None,
32            label: None,
33        }
34    }
35
36    /// Set the management-call sender used for `install_canister`.
37    #[must_use]
38    pub const fn install_sender(mut self, sender: Principal) -> Self {
39        self.install_sender = Some(sender);
40        self
41    }
42
43    /// Set a diagnostic label for install failures.
44    #[must_use]
45    pub fn label(mut self, label: impl Into<String>) -> Self {
46        self.label = Some(label.into());
47        self
48    }
49}
50
51/// Retry limits and simulated cooldown for install-code operations.
52#[derive(Clone, Copy, Debug, Eq, PartialEq)]
53pub struct RetryPolicy {
54    max_attempts: NonZeroUsize,
55    cooldown: Duration,
56}
57
58impl RetryPolicy {
59    /// Create a policy with an exact, non-zero maximum attempt count.
60    #[must_use]
61    pub const fn new(max_attempts: NonZeroUsize, cooldown: Duration) -> Self {
62        Self {
63            max_attempts,
64            cooldown,
65        }
66    }
67
68    /// Read the maximum number of operation attempts, including the first.
69    #[must_use]
70    pub const fn max_attempts(self) -> NonZeroUsize {
71        self.max_attempts
72    }
73
74    /// Read the simulated cooldown applied between rate-limited attempts.
75    #[must_use]
76    pub const fn cooldown(self) -> Duration {
77        self.cooldown
78    }
79}
80
81/// Generic canister installation and structured install-code retry policy.
82///
83/// The extension creates canisters using PocketIC defaults. `InstallSpec::cycles`
84/// is an additional top-up, not the complete initial balance.
85pub trait CanisterInstallExt {
86    /// Create and install one canister from a reusable specification.
87    ///
88    /// Captures PocketIC panics during creation, cycle funding, and installation.
89    /// The error identifies the stage and retains an id when creation succeeded.
90    /// Callers choose whether to collect diagnostics; this performs no diagnostic
91    /// requests or error logging.
92    fn try_create_and_install(&self, spec: InstallSpec) -> Result<Principal, CanisterInstallError>;
93
94    /// Sequentially create and install multiple canisters, stopping at the first failure.
95    fn try_create_and_install_many<I>(
96        &self,
97        specs: I,
98    ) -> Result<Vec<Principal>, CanisterInstallError>
99    where
100        I: IntoIterator<Item = InstallSpec>;
101
102    /// Advance simulated time once, then tick twice for an install-code cooldown.
103    /// Zero duration skips time reads/writes while retaining both ticks.
104    fn wait_out_install_code_rate_limit(&self, cooldown: Duration);
105
106    /// Retry only while PocketIC reports install-code rate limiting.
107    ///
108    /// `RetryPolicy::max_attempts` includes the initial call. Before each
109    /// retry, this advances simulated time by the configured cooldown and
110    /// executes two PocketIC ticks. Other rejections are returned unchanged.
111    fn retry_install_code<T, F>(&self, policy: RetryPolicy, op: F) -> Result<T, RejectResponse>
112    where
113        F: FnMut() -> Result<T, RejectResponse>;
114}
115
116impl CanisterInstallExt for PocketIc {
117    /// Install one arbitrary wasm module from a generic install specification.
118    fn try_create_and_install(&self, spec: InstallSpec) -> Result<Principal, CanisterInstallError> {
119        let canister_id = capture_operation(|| self.create_canister()).map_err(|source| {
120            CanisterInstallError::CreateCanister {
121                label: spec.label.clone(),
122                source,
123            }
124        })?;
125        if spec.cycles > 0 {
126            capture_operation(|| self.add_cycles(canister_id, spec.cycles)).map_err(|source| {
127                CanisterInstallError::AddCycles {
128                    canister_id,
129                    label: spec.label.clone(),
130                    source,
131                }
132            })?;
133        }
134
135        // Move the Wasm and init payload while retaining the label for failures.
136        let install = capture_operation(|| {
137            self.install_canister(canister_id, spec.wasm, spec.init_bytes, spec.install_sender);
138        });
139        if let Err(source) = install {
140            return Err(CanisterInstallError::InstallCode {
141                canister_id,
142                label: spec.label,
143                source,
144            });
145        }
146
147        Ok(canister_id)
148    }
149
150    /// Sequentially install multiple arbitrary wasm modules into this PocketIC instance.
151    ///
152    /// Installs are attempted in iterator order. If one install fails, earlier
153    /// installs remain in the PocketIC instance, the failed canister may exist
154    /// with the id retained in the funding or installation error variant, and later
155    /// installs are not attempted.
156    fn try_create_and_install_many<I>(
157        &self,
158        specs: I,
159    ) -> Result<Vec<Principal>, CanisterInstallError>
160    where
161        I: IntoIterator<Item = InstallSpec>,
162    {
163        specs
164            .into_iter()
165            .map(|spec| self.try_create_and_install(spec))
166            .collect()
167    }
168
169    /// Wait out the PocketIC `install_code` cooldown window inside the same instance.
170    fn wait_out_install_code_rate_limit(&self, cooldown: Duration) {
171        super::time::advance_and_tick(self, cooldown, 2);
172    }
173
174    fn retry_install_code<T, F>(&self, policy: RetryPolicy, op: F) -> Result<T, RejectResponse>
175    where
176        F: FnMut() -> Result<T, RejectResponse>,
177    {
178        retry_install_code_with(policy, op, || {
179            self.wait_out_install_code_rate_limit(policy.cooldown());
180        })
181    }
182}
183
184fn retry_install_code_with<T, F, W>(
185    policy: RetryPolicy,
186    mut op: F,
187    mut wait_out_cooldown: W,
188) -> Result<T, RejectResponse>
189where
190    F: FnMut() -> Result<T, RejectResponse>,
191    W: FnMut(),
192{
193    let mut retries = 1..policy.max_attempts().get();
194    loop {
195        match op() {
196            Err(err)
197                if err.error_code == ErrorCode::CanisterInstallCodeRateLimited
198                    && retries.next().is_some() =>
199            {
200                wait_out_cooldown();
201            }
202            result => return result,
203        }
204    }
205}
206
207#[cfg(test)]
208mod tests {
209    use std::{
210        cell::{Cell, RefCell},
211        num::NonZeroUsize,
212        time::Duration,
213    };
214
215    use pocket_ic::{ErrorCode, RejectCode, RejectResponse};
216
217    use super::{RetryPolicy, retry_install_code_with};
218
219    fn rejection(error_code: ErrorCode, message: &str) -> RejectResponse {
220        RejectResponse {
221            reject_code: RejectCode::SysTransient,
222            reject_message: message.to_string(),
223            error_code,
224            certified: false,
225        }
226    }
227
228    #[test]
229    fn retry_policy_counts_the_first_attempt() {
230        for max_attempts in [1, 3] {
231            let attempts = Cell::new(0);
232            let waits = Cell::new(0);
233            let result = retry_install_code_with(
234                RetryPolicy::new(
235                    NonZeroUsize::new(max_attempts).unwrap(),
236                    Duration::from_secs(1),
237                ),
238                || {
239                    attempts.set(attempts.get() + 1);
240                    Err::<(), _>(rejection(
241                        ErrorCode::CanisterInstallCodeRateLimited,
242                        &format!("rate limit on attempt {}", attempts.get()),
243                    ))
244                },
245                || waits.set(waits.get() + 1),
246            );
247
248            assert_eq!(
249                result,
250                Err(rejection(
251                    ErrorCode::CanisterInstallCodeRateLimited,
252                    &format!("rate limit on attempt {max_attempts}"),
253                ))
254            );
255            assert_eq!(attempts.get(), max_attempts);
256            assert_eq!(waits.get(), max_attempts - 1);
257        }
258    }
259
260    #[test]
261    fn retry_policy_waits_only_between_retryable_attempts() {
262        for success_attempt in [1, 3] {
263            let events = RefCell::new(Vec::new());
264            let attempts = Cell::new(0);
265            let result = retry_install_code_with(
266                RetryPolicy::new(NonZeroUsize::new(3).unwrap(), Duration::from_secs(1)),
267                || {
268                    events.borrow_mut().push("attempt");
269                    attempts.set(attempts.get() + 1);
270                    if attempts.get() == success_attempt {
271                        Ok(42)
272                    } else {
273                        Err(rejection(
274                            ErrorCode::CanisterInstallCodeRateLimited,
275                            "rate limit",
276                        ))
277                    }
278                },
279                || events.borrow_mut().push("wait"),
280            );
281            assert_eq!(result, Ok(42));
282            assert_eq!(attempts.get(), success_attempt);
283            assert_eq!(
284                events.into_inner(),
285                if success_attempt == 1 {
286                    vec!["attempt"]
287                } else {
288                    vec!["attempt", "wait", "attempt", "wait", "attempt"]
289                }
290            );
291        }
292    }
293
294    #[test]
295    fn retry_policy_stops_on_non_rate_limit_failure() {
296        let attempts = Cell::new(0);
297        let not_retryable = rejection(ErrorCode::CanisterRejectedMessage, "not retryable");
298        let result = retry_install_code_with(
299            RetryPolicy::new(NonZeroUsize::new(3).unwrap(), Duration::from_secs(1)),
300            || {
301                attempts.set(attempts.get() + 1);
302                Err::<(), _>(not_retryable.clone())
303            },
304            || panic!("non-rate-limit failure must not wait"),
305        );
306
307        assert_eq!(result, Err(not_retryable));
308        assert_eq!(attempts.get(), 1);
309    }
310}
ic_testkit/pic/lifecycle.rs

ic_testkit/pic/
lifecycle.rs
