ic_testkit/pic/
baseline.rs1use super::{
2 CanisterRestoreReceipt, CanisterSnapshotTarget, ControllerSnapshotError, ControllerSnapshots,
3 PocketIcInstance, 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(
46 &self,
47 funding: SnapshotRestoreFunding,
48 ) -> Result<CanisterRestoreReceipt, ControllerSnapshotError> {
49 self.instance
50 .instance
51 .restore_snapshots(&self.snapshots, funding)?;
52 Ok(CanisterRestoreReceipt::from_restore(
53 &self.snapshots,
54 funding,
55 ))
56 }
57
58 #[must_use]
60 pub const fn pocket_ic(&self) -> &PocketIc {
61 &self.instance.instance
62 }
63
64 #[must_use]
66 pub fn snapshot_count(&self) -> usize {
67 self.snapshots.len()
68 }
69
70 pub fn snapshot_canister_ids(&self) -> impl Iterator<Item = Principal> + '_ {
72 self.snapshots.canister_ids()
73 }
74
75 #[must_use]
77 pub const fn metadata(&self) -> &T {
78 &self.metadata
79 }
80
81 #[must_use]
83 pub const fn metadata_mut(&mut self) -> &mut T {
84 &mut self.metadata
85 }
86}