ic_testkit/pic/
lifecycle.rs

1use std::{num::NonZeroUsize, time::Duration};
2
3use ic_platform::pocket_ic::{ErrorCode, PocketIc, RejectResponse};
4use ic_platform::principal::ExternalPrincipal as Principal;
5
6use super::{CanisterInstallError, transport::capture_operation};
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    /// Captures `PocketIC` panics 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    fn wait_out_install_code_rate_limit(&self, cooldown: Duration);
50
51    /// Retry only while `PocketIC` reports install-code rate limiting.
52    ///
53    /// `RetryPolicy::max_attempts` includes the initial call. Before each
54    /// retry, this advances simulated time by the configured cooldown and
55    /// executes two `PocketIC` ticks. Other rejections are returned unchanged.
56    ///
57    /// # Errors
58    /// Returns the final rejection after the selected retry policy is exhausted.
59    fn retry_install_code<T, F>(&self, policy: RetryPolicy, op: F) -> Result<T, RejectResponse>
60    where
61        F: FnMut() -> Result<T, RejectResponse>;
62}
63
64impl CanisterInstallExt for PocketIc {
65    /// Install one arbitrary wasm module from a generic install specification.
66    fn try_create_and_install(&self, spec: InstallSpec) -> Result<Principal, CanisterInstallError> {
67        let canister_id = capture_operation(|| self.create_canister()).map_err(|source| {
68            CanisterInstallError::CreateCanister {
69                label: spec.label.clone(),
70                source,
71            }
72        })?;
73        if spec.cycles > 0 {
74            capture_operation(|| self.add_cycles(canister_id, spec.cycles)).map_err(|source| {
75                CanisterInstallError::AddCycles {
76                    canister_id,
77                    label: spec.label.clone(),
78                    source,
79                }
80            })?;
81        }
82
83        // Move the Wasm and init payload while retaining the label for failures.
84        let install = capture_operation(|| {
85            self.install_canister(canister_id, spec.wasm, spec.init_bytes, spec.install_sender);
86        });
87        if let Err(source) = install {
88            return Err(CanisterInstallError::InstallCode {
89                canister_id,
90                label: spec.label,
91                source,
92            });
93        }
94
95        Ok(canister_id)
96    }
97
98    /// Wait out the `PocketIC` `install_code` cooldown window inside the same instance.
99    fn wait_out_install_code_rate_limit(&self, cooldown: Duration) {
100        super::time::advance_and_tick(self, cooldown, 2);
101    }
102
103    fn retry_install_code<T, F>(&self, policy: RetryPolicy, op: F) -> Result<T, RejectResponse>
104    where
105        F: FnMut() -> Result<T, RejectResponse>,
106    {
107        retry_install_code_with(policy, op, || {
108            self.wait_out_install_code_rate_limit(policy.cooldown);
109        })
110    }
111}
112
113fn retry_install_code_with<T, F, W>(
114    policy: RetryPolicy,
115    mut op: F,
116    mut wait_out_cooldown: W,
117) -> Result<T, RejectResponse>
118where
119    F: FnMut() -> Result<T, RejectResponse>,
120    W: FnMut(),
121{
122    let mut retries = 1..policy.max_attempts.get();
123    loop {
124        match op() {
125            Err(err)
126                if err.error_code == ErrorCode::CanisterInstallCodeRateLimited
127                    && retries.next().is_some() =>
128            {
129                wait_out_cooldown();
130            }
131            result => return result,
132        }
133    }
134}
135
136#[cfg(test)]
137mod tests {
138    use std::{
139        cell::{Cell, RefCell},
140        num::NonZeroUsize,
141        time::Duration,
142    };
143
144    use ic_platform::pocket_ic::{ErrorCode, RejectCode, RejectResponse};
145
146    use super::{RetryPolicy, retry_install_code_with};
147
148    fn rejection(error_code: ErrorCode, message: &str) -> RejectResponse {
149        RejectResponse {
150            reject_code: RejectCode::SysTransient,
151            reject_message: message.to_string(),
152            error_code,
153            certified: false,
154        }
155    }
156
157    #[test]
158    fn retry_policy_counts_the_first_attempt() {
159        for max_attempts in [1, 3] {
160            let attempts = Cell::new(0);
161            let waits = Cell::new(0);
162            let result = retry_install_code_with(
163                RetryPolicy {
164                    max_attempts: NonZeroUsize::new(max_attempts).unwrap(),
165                    cooldown: Duration::from_secs(1),
166                },
167                || {
168                    attempts.set(attempts.get() + 1);
169                    Err::<(), _>(rejection(
170                        ErrorCode::CanisterInstallCodeRateLimited,
171                        &format!("rate limit on attempt {}", attempts.get()),
172                    ))
173                },
174                || waits.set(waits.get() + 1),
175            );
176
177            assert_eq!(
178                result,
179                Err(rejection(
180                    ErrorCode::CanisterInstallCodeRateLimited,
181                    &format!("rate limit on attempt {max_attempts}"),
182                ))
183            );
184            assert_eq!(attempts.get(), max_attempts);
185            assert_eq!(waits.get(), max_attempts - 1);
186        }
187    }
188
189    #[test]
190    fn retry_policy_waits_only_between_retryable_attempts() {
191        for success_attempt in [1, 3] {
192            let events = RefCell::new(Vec::new());
193            let attempts = Cell::new(0);
194            let result = retry_install_code_with(
195                RetryPolicy {
196                    max_attempts: NonZeroUsize::new(3).unwrap(),
197                    cooldown: Duration::from_secs(1),
198                },
199                || {
200                    events.borrow_mut().push("attempt");
201                    attempts.set(attempts.get() + 1);
202                    if attempts.get() == success_attempt {
203                        Ok(42)
204                    } else {
205                        Err(rejection(
206                            ErrorCode::CanisterInstallCodeRateLimited,
207                            "rate limit",
208                        ))
209                    }
210                },
211                || events.borrow_mut().push("wait"),
212            );
213            assert_eq!(result, Ok(42));
214            assert_eq!(attempts.get(), success_attempt);
215            assert_eq!(
216                events.into_inner(),
217                if success_attempt == 1 {
218                    vec!["attempt"]
219                } else {
220                    vec!["attempt", "wait", "attempt", "wait", "attempt"]
221                }
222            );
223        }
224    }
225
226    #[test]
227    fn retry_policy_stops_on_non_rate_limit_failure() {
228        let attempts = Cell::new(0);
229        let not_retryable = rejection(ErrorCode::CanisterRejectedMessage, "not retryable");
230        let result = retry_install_code_with(
231            RetryPolicy {
232                max_attempts: NonZeroUsize::new(3).unwrap(),
233                cooldown: Duration::from_secs(1),
234            },
235            || {
236                attempts.set(attempts.get() + 1);
237                Err::<(), _>(not_retryable.clone())
238            },
239            || panic!("non-rate-limit failure must not wait"),
240        );
241
242        assert_eq!(result, Err(not_retryable));
243        assert_eq!(attempts.get(), 1);
244    }
245}
ic_testkit/pic/lifecycle.rs

ic_testkit/pic/
lifecycle.rs
