ic_testkit/pic/baseline.rs
1use super::{
2 CanisterSnapshotTarget, ControllerSnapshotError, ControllerSnapshots, PocketIcInstance,
3 PocketIcSnapshotExt, SnapshotRestoreFunding,
4};
5use crate::pocket_ic::PocketIc;
6use ic_platform_types::principal::ExternalPrincipal as Principal;
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 a cached baseline with an explicit sender for every canister.
20 ///
21 /// Empty inputs are rejected before capture. Capture is ordered and transactional as documented by
22 /// [`PocketIcSnapshotExt::capture_snapshots`]. No sender fallback is performed.
23 /// The caller stops/drains every target before capture and restarts only
24 /// after the entire capture succeeds. Automatic progress must be disabled.
25 ///
26 /// # Errors
27 /// Returns a controller snapshot error if target admission or snapshot capture fails.
28 pub async fn capture<I>(
29 instance: PocketIcInstance<PocketIc>,
30 targets: I,
31 metadata: T,
32 ) -> Result<Self, ControllerSnapshotError>
33 where
34 I: IntoIterator<Item = CanisterSnapshotTarget>,
35 {
36 let mut targets = targets.into_iter().peekable();
37 let capture = if targets.peek().is_none() {
38 Err(ControllerSnapshotError::EmptyCanisterSet)
39 } else {
40 instance.instance().capture_snapshots(targets).await
41 };
42 let snapshots = match capture {
43 Ok(snapshots) => snapshots,
44 Err(error) => {
45 if let Err(cleanup) = instance.close().await {
46 tracing::warn!("PocketIC deletion after capture failure failed: {cleanup}");
47 }
48 return Err(error);
49 }
50 };
51 Ok(Self {
52 instance,
53 snapshots,
54 metadata,
55 })
56 }
57
58 /// Restore with exactly the captured senders and selected cycle funding.
59 /// Success means the complete captured set was restored.
60 /// Stop/drain the entire set first; restart after complete restore/reset.
61 /// Absolute captured timers and hook state are restored; current time is not.
62 ///
63 /// # Errors
64 /// Returns a controller snapshot error if funding or snapshot restoration fails.
65 #[expect(
66 clippy::future_not_send,
67 reason = "baseline metadata may be local; pool recipes separately require Sync metadata"
68 )]
69 pub async fn restore(
70 &self,
71 funding: SnapshotRestoreFunding,
72 ) -> Result<(), ControllerSnapshotError> {
73 self.instance
74 .instance()
75 .restore_snapshots(&self.snapshots, funding)
76 .await
77 }
78
79 /// Explicitly delete the instance before releasing its managed server.
80 /// Cancellation on a borrowed server requires caller reconciliation.
81 /// # Errors
82 /// Returns the client deletion failure, without retrying it.
83 pub async fn close(self) -> Result<(), crate::pocket_ic::ClientError> {
84 self.instance.close().await
85 }
86
87 /// Borrow the owned `PocketIC` instance behind this cached baseline.
88 #[must_use]
89 pub const fn pocket_ic(&self) -> &PocketIc {
90 self.instance.instance()
91 }
92
93 /// Return the number of canisters captured by this baseline.
94 #[must_use]
95 pub fn snapshot_count(&self) -> usize {
96 self.snapshots.len()
97 }
98
99 /// Iterate over captured canister ids in deterministic principal order.
100 pub fn snapshot_canister_ids(&self) -> impl Iterator<Item = Principal> + '_ {
101 self.snapshots.canister_ids()
102 }
103
104 /// Borrow the captured metadata associated with this cached baseline.
105 #[must_use]
106 pub const fn metadata(&self) -> &T {
107 &self.metadata
108 }
109
110 /// Mutably borrow the captured metadata associated with this cached baseline.
111 #[must_use]
112 pub const fn metadata_mut(&mut self) -> &mut T {
113 &mut self.metadata
114 }
115}