ic_testkit/pic/
baseline.rs1use super::{
2 CanisterSnapshotTarget, ControllerSnapshotError, ControllerSnapshots, PocketIcInstance,
3 PocketIcSnapshotExt, SnapshotRestoreFunding,
4};
5use candid::Principal;
6use pocket_ic::PocketIc;
7
8pub struct CachedPocketIcBaseline<T> {
13 instance: PocketIcInstance<PocketIc>,
14 snapshots: ControllerSnapshots,
15 metadata: T,
16}
17
18impl<T> CachedPocketIcBaseline<T> {
19 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 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 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 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 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 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 #[must_use]
100 pub const fn pocket_ic(&self) -> &PocketIc {
101 &self.instance.instance
102 }
103
104 #[must_use]
106 pub fn snapshot_count(&self) -> usize {
107 self.snapshots.len()
108 }
109
110 pub fn snapshot_canister_ids(&self) -> impl Iterator<Item = Principal> + '_ {
112 self.snapshots.canister_ids()
113 }
114
115 #[must_use]
117 pub const fn metadata(&self) -> &T {
118 &self.metadata
119 }
120
121 #[must_use]
123 pub const fn metadata_mut(&mut self) -> &mut T {
124 &mut self.metadata
125 }
126}