ic_testkit/pic/
lifecycle.rs

1use std::panic::{AssertUnwindSafe, catch_unwind};
2use std::time::Duration;
3
4use candid::Principal;
5
6use super::{Pic, PicInstallError, startup};
7
8impl Pic {
9    /// Install one arbitrary wasm module with caller-provided init bytes.
10    ///
11    /// This is the generic install path for downstreams that use `ic-testkit`
12    /// without depending on application-specific init payload conventions.
13    #[must_use]
14    pub fn create_and_install_with_args(
15        &self,
16        wasm: Vec<u8>,
17        init_bytes: Vec<u8>,
18        install_cycles: u128,
19    ) -> Principal {
20        self.try_create_and_install_with_args(wasm, init_bytes, install_cycles)
21            .unwrap_or_else(|err| panic!("{err}"))
22    }
23
24    /// Install one arbitrary wasm module with caller-provided init bytes.
25    pub fn try_create_and_install_with_args(
26        &self,
27        wasm: Vec<u8>,
28        init_bytes: Vec<u8>,
29        install_cycles: u128,
30    ) -> Result<Principal, PicInstallError> {
31        self.try_create_funded_and_install(wasm, init_bytes, install_cycles)
32    }
33
34    /// Wait out the PocketIC `install_code` cooldown window inside the same instance.
35    pub fn wait_out_install_code_rate_limit(&self, cooldown: Duration) {
36        self.advance_time(cooldown);
37        self.tick_n(2);
38    }
39
40    /// Retry one install_code-like operation while PocketIC still reports rate limiting.
41    pub fn retry_install_code_ok<T, F>(
42        &self,
43        retry_limit: usize,
44        cooldown: Duration,
45        mut op: F,
46    ) -> Result<T, String>
47    where
48        F: FnMut() -> Result<T, String>,
49    {
50        let mut last_err = None;
51
52        for _ in 0..retry_limit {
53            match op() {
54                Ok(value) => return Ok(value),
55                Err(err) if is_install_code_rate_limited(&err) => {
56                    last_err = Some(err);
57                    self.wait_out_install_code_rate_limit(cooldown);
58                }
59                Err(err) => return Err(err),
60            }
61        }
62
63        Err(last_err.unwrap_or_else(|| "install_code retry loop exhausted".to_string()))
64    }
65
66    /// Retry one install_code-like failure path while PocketIC still reports rate limiting.
67    pub fn retry_install_code_err<F>(
68        &self,
69        retry_limit: usize,
70        cooldown: Duration,
71        first: Result<(), String>,
72        mut op: F,
73    ) -> Result<(), String>
74    where
75        F: FnMut() -> Result<(), String>,
76    {
77        match first {
78            Ok(()) => return Ok(()),
79            Err(err) if !is_install_code_rate_limited(&err) => return Err(err),
80            Err(_) => {}
81        }
82
83        self.wait_out_install_code_rate_limit(cooldown);
84
85        for _ in 1..retry_limit {
86            match op() {
87                Ok(()) => return Ok(()),
88                Err(err) if is_install_code_rate_limited(&err) => {
89                    self.wait_out_install_code_rate_limit(cooldown);
90                }
91                Err(err) => return Err(err),
92            }
93        }
94
95        op()
96    }
97
98    // Install a canister after creating it and funding it with cycles.
99    fn try_create_funded_and_install(
100        &self,
101        wasm: Vec<u8>,
102        init_bytes: Vec<u8>,
103        install_cycles: u128,
104    ) -> Result<Principal, PicInstallError> {
105        let canister_id = self.create_canister();
106        self.add_cycles(canister_id, install_cycles);
107
108        let install = catch_unwind(AssertUnwindSafe(|| {
109            self.inner
110                .install_canister(canister_id, wasm, init_bytes, None);
111        }));
112        if let Err(payload) = install {
113            eprintln!("install_canister trapped for {canister_id}");
114            if let Ok(status) = self.inner.canister_status(canister_id, None) {
115                eprintln!("canister_status for {canister_id}: {status:?}");
116            }
117            if let Ok(logs) = self
118                .inner
119                .fetch_canister_logs(canister_id, Principal::anonymous())
120            {
121                for record in logs {
122                    eprintln!("canister_log {canister_id}: {record:?}");
123                }
124            }
125            return Err(PicInstallError::new(
126                canister_id,
127                startup::panic_payload_to_string(payload.as_ref()),
128            ));
129        }
130
131        Ok(canister_id)
132    }
133}
134
135fn is_install_code_rate_limited(message: &str) -> bool {
136    message.contains("CanisterInstallCodeRateLimited")
137}
ic_testkit/pic/lifecycle.rs

ic_testkit/pic/
lifecycle.rs
