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 targets: I,
26 metadata: T,
27 ) -> Result<Self, ControllerSnapshotError>
28 where
29 I: IntoIterator<Item = CanisterSnapshotTarget>,
30 {
31 let mut targets = targets.into_iter().peekable();
32 if targets.peek().is_none() {
33 return Err(ControllerSnapshotError::EmptyCanisterSet);
34 }
35 let snapshots = instance.instance.capture_snapshots(targets)?;
36 Ok(Self {
37 instance,
38 snapshots,
39 metadata,
40 })
41 }
42
43 pub fn restore(&self, funding: SnapshotRestoreFunding) -> Result<(), ControllerSnapshotError> {
46 self.instance
47 .instance
48 .restore_snapshots(&self.snapshots, funding)
49 }
50
51 #[must_use]
53 pub const fn pocket_ic(&self) -> &PocketIc {
54 &self.instance.instance
55 }
56
57 #[must_use]
59 pub fn snapshot_count(&self) -> usize {
60 self.snapshots.len()
61 }
62
63 pub fn snapshot_canister_ids(&self) -> impl Iterator<Item = Principal> + '_ {
65 self.snapshots.canister_ids()
66 }
67
68 #[must_use]
70 pub const fn metadata(&self) -> &T {
71 &self.metadata
72 }
73
74 #[must_use]
76 pub const fn metadata_mut(&mut self) -> &mut T {
77 &mut self.metadata
78 }
79}