Skip to main content

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}