Accept custom snapshot sources in Env::from_ledger_snapshot

soroban-sdk/src/env.rs

@@ -482,7 +482,7 @@ use crate::{
     testutils::{
         budget::Budget, default_ledger_info, Address as _, AuthSnapshot, AuthorizedInvocation,
         ContractFunctionSet, EventsSnapshot, Generators, Ledger as _, MockAuth, MockAuthContract,
-        Register, Snapshot, StellarAssetContract, StellarAssetIssuer,
+        Register, Snapshot, SnapshotSourceInput, StellarAssetContract, StellarAssetIssuer,
     },
     Bytes, BytesN, ConstructorArgs,
 };
@@ -1600,16 +1600,22 @@ impl Env {
         self.to_snapshot().write_file(p).unwrap();
     }
 
-    /// Creates a new Env loaded with the [`LedgerSnapshot`].
+    /// Creates a new Env loaded with the snapshot source.
     ///
-    /// The ledger info and state in the snapshot are loaded into the Env.
-    pub fn from_ledger_snapshot(s: LedgerSnapshot) -> Env {
+    /// The ledger info and state from the snapshot source are loaded into the Env.
+    pub fn from_ledger_snapshot(input: impl Into<SnapshotSourceInput>) -> Env {
+        let SnapshotSourceInput {
+            source,
+            ledger_info,
+            snapshot,
+        } = input.into();
+
         Env::new_for_testutils(
             EnvTestConfig::default(), // TODO: Allow setting the config.
-            Rc::new(s.clone()),
+            source,
             None,
-            Some(s.ledger_info()),
-            Some(Rc::new(s.clone())),
+            ledger_info,
+            snapshot,
         )
     }
 

soroban-sdk/src/testutils.rs

@@ -25,6 +25,8 @@ use soroban_ledger_snapshot::LedgerSnapshot;
 
 pub use crate::env::EnvTestConfig;
 
+pub use crate::env::internal::{storage::SnapshotSource, HostError};
+
 pub trait Register {
     fn register<'i, I, A>(self, env: &Env, id: I, args: A) -> crate::Address
     where
@@ -605,3 +607,20 @@ impl StellarAssetContract {
         self.asset.clone()
     }
 }
+
+pub struct SnapshotSourceInput {
+    pub source: Rc<dyn SnapshotSource>,
+    pub ledger_info: Option<LedgerInfo>,
+    pub snapshot: Option<Rc<LedgerSnapshot>>,
+}
+
+impl From<LedgerSnapshot> for SnapshotSourceInput {
+    fn from(s: LedgerSnapshot) -> Self {
+        let s = Rc::new(s);
+        Self {
+            source: s.clone(),
+            ledger_info: Some(s.ledger_info()),
+            snapshot: Some(s),
+        }
+    }
+}