ic_testkit/pic/
baseline.rs

1use super::{
2    CanisterSnapshotTarget, ControllerSnapshotError, ControllerSnapshots, PocketIcSnapshotExt,
3    SnapshotRestoreFunding,
4};
5use candid::Principal;
6use pocket_ic::PocketIc;
7
8/// One owned PocketIC instance with captured snapshots and caller metadata.
9///
10/// The value contains no synchronization. A [`super::PocketIcBaselineRecipe`]
11/// can own it as its fixture in a [`super::CachedPocketIcBaselinePool`].
12pub struct CachedPocketIcBaseline<T> {
13    pocket_ic: PocketIc,
14    snapshots: ControllerSnapshots,
15    metadata: T,
16}
17
18impl<T> CachedPocketIcBaseline<T> {
19    /// Capture one cached baseline from the current PocketIC instance.
20    ///
21    /// Snapshot capture is ordered and transactional as documented by
22    /// [`PocketIcSnapshotExt::capture_controller_snapshots`].
23    pub fn capture<I>(
24        pocket_ic: PocketIc,
25        controller_id: Principal,
26        canister_ids: I,
27        metadata: T,
28    ) -> Result<Self, ControllerSnapshotError>
29    where
30        I: IntoIterator<Item = Principal>,
31    {
32        let snapshots = pocket_ic.capture_controller_snapshots(controller_id, canister_ids)?;
33
34        Ok(Self {
35            pocket_ic,
36            snapshots,
37            metadata,
38        })
39    }
40
41    /// Capture one cached baseline with an explicit sender for every canister.
42    ///
43    /// This avoids fallback rejections in mixed-controller topologies.
44    pub fn capture_with_senders<I>(
45        pocket_ic: PocketIc,
46        targets: I,
47        metadata: T,
48    ) -> Result<Self, ControllerSnapshotError>
49    where
50        I: IntoIterator<Item = CanisterSnapshotTarget>,
51    {
52        let snapshots = pocket_ic.capture_snapshots_with_senders(targets)?;
53
54        Ok(Self {
55            pocket_ic,
56            snapshots,
57            metadata,
58        })
59    }
60
61    /// Restore the captured snapshot set without adding cycles.
62    pub fn restore(&self, controller_id: Principal) -> Result<(), ControllerSnapshotError> {
63        self.pocket_ic
64            .restore_controller_snapshots(controller_id, &self.snapshots)
65    }
66
67    /// Restore the captured snapshot set with an explicit cycle-funding policy.
68    pub fn restore_with_funding(
69        &self,
70        controller_id: Principal,
71        funding: SnapshotRestoreFunding,
72    ) -> Result<(), ControllerSnapshotError> {
73        self.pocket_ic.restore_controller_snapshots_with_funding(
74            controller_id,
75            &self.snapshots,
76            funding,
77        )
78    }
79
80    /// Restore with exactly the senders retained during snapshot capture.
81    pub fn restore_with_captured_senders(&self) -> Result<(), ControllerSnapshotError> {
82        self.pocket_ic
83            .restore_snapshots_with_captured_senders(&self.snapshots)
84    }
85
86    /// Restore with captured senders and an explicit cycle-funding policy.
87    pub fn restore_with_captured_senders_and_funding(
88        &self,
89        funding: SnapshotRestoreFunding,
90    ) -> Result<(), ControllerSnapshotError> {
91        self.pocket_ic
92            .restore_snapshots_with_captured_senders_and_funding(&self.snapshots, funding)
93    }
94
95    /// Borrow the owned PocketIC instance behind this cached baseline.
96    #[must_use]
97    pub const fn pocket_ic(&self) -> &PocketIc {
98        &self.pocket_ic
99    }
100
101    /// Return the number of canisters captured by this baseline.
102    #[must_use]
103    pub fn snapshot_count(&self) -> usize {
104        self.snapshots.len()
105    }
106
107    /// Iterate over captured canister ids in deterministic principal order.
108    pub fn snapshot_canister_ids(&self) -> impl Iterator<Item = Principal> + '_ {
109        self.snapshots.canister_ids()
110    }
111
112    /// Borrow the captured metadata associated with this cached baseline.
113    #[must_use]
114    pub const fn metadata(&self) -> &T {
115        &self.metadata
116    }
117
118    /// Mutably borrow the captured metadata associated with this cached baseline.
119    #[must_use]
120    pub const fn metadata_mut(&mut self) -> &mut T {
121        &mut self.metadata
122    }
123}
ic_testkit/pic/baseline.rs

ic_testkit/pic/
baseline.rs
