Skip to main content

ic_testkit/pic/
baseline.rs

1use super::{
2    CanisterSnapshotTarget, ControllerSnapshotError, ControllerSnapshots, PocketIcInstance,
3    PocketIcSnapshotExt, SnapshotRestoreFunding,
4};
5use candid::Principal;
6use pocket_ic::PocketIc;
7
8/// One owned PocketIC instance/server record with 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    instance: PocketIcInstance<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        instance: PocketIcInstance<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 = instance
33            .instance
34            .capture_controller_snapshots(controller_id, canister_ids)?;
35
36        Ok(Self {
37            instance,
38            snapshots,
39            metadata,
40        })
41    }
42
43    /// Capture one cached baseline with an explicit sender for every canister.
44    ///
45    /// This avoids fallback rejections in mixed-controller topologies.
46    pub fn capture_with_senders<I>(
47        instance: PocketIcInstance<PocketIc>,
48        targets: I,
49        metadata: T,
50    ) -> Result<Self, ControllerSnapshotError>
51    where
52        I: IntoIterator<Item = CanisterSnapshotTarget>,
53    {
54        let snapshots = instance.instance.capture_snapshots_with_senders(targets)?;
55
56        Ok(Self {
57            instance,
58            snapshots,
59            metadata,
60        })
61    }
62
63    /// Restore the captured snapshot set without adding cycles.
64    pub fn restore(&self, controller_id: Principal) -> Result<(), ControllerSnapshotError> {
65        self.instance
66            .instance
67            .restore_controller_snapshots(controller_id, &self.snapshots)
68    }
69
70    /// Restore the captured snapshot set with an explicit cycle-funding policy.
71    pub fn restore_with_funding(
72        &self,
73        controller_id: Principal,
74        funding: SnapshotRestoreFunding,
75    ) -> Result<(), ControllerSnapshotError> {
76        self.instance
77            .instance
78            .restore_controller_snapshots_with_funding(controller_id, &self.snapshots, funding)
79    }
80
81    /// Restore with exactly the senders retained during snapshot capture.
82    pub fn restore_with_captured_senders(&self) -> Result<(), ControllerSnapshotError> {
83        self.instance
84            .instance
85            .restore_snapshots_with_captured_senders(&self.snapshots)
86    }
87
88    /// Restore with captured senders and an explicit cycle-funding policy.
89    pub fn restore_with_captured_senders_and_funding(
90        &self,
91        funding: SnapshotRestoreFunding,
92    ) -> Result<(), ControllerSnapshotError> {
93        self.instance
94            .instance
95            .restore_snapshots_with_captured_senders_and_funding(&self.snapshots, funding)
96    }
97
98    /// Borrow the owned PocketIC instance behind this cached baseline.
99    #[must_use]
100    pub const fn pocket_ic(&self) -> &PocketIc {
101        &self.instance.instance
102    }
103
104    /// Return the number of canisters captured by this baseline.
105    #[must_use]
106    pub fn snapshot_count(&self) -> usize {
107        self.snapshots.len()
108    }
109
110    /// Iterate over captured canister ids in deterministic principal order.
111    pub fn snapshot_canister_ids(&self) -> impl Iterator<Item = Principal> + '_ {
112        self.snapshots.canister_ids()
113    }
114
115    /// Borrow the captured metadata associated with this cached baseline.
116    #[must_use]
117    pub const fn metadata(&self) -> &T {
118        &self.metadata
119    }
120
121    /// Mutably borrow the captured metadata associated with this cached baseline.
122    #[must_use]
123    pub const fn metadata_mut(&mut self) -> &mut T {
124        &mut self.metadata
125    }
126}