feat: only use snapshots if ENV present

Author: hstove

testnet/stacks-node/src/tests/signer/mod.rs

@@ -19,8 +19,8 @@ use std::collections::HashSet;
 use std::path::PathBuf;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{Arc, Mutex};
-use std::thread;
 use std::time::{Duration, Instant};
+use std::{env, thread};
 
 use clarity::boot_util::boot_code_id;
 use clarity::vm::types::PrincipalData;
@@ -110,6 +110,16 @@ pub struct SignerTest<S> {
     pub snapshot_path: Option<PathBuf>,
 }
 
+struct SnapshotSetupInfo {
+    snapshot_path: PathBuf,
+    snapshot_exists: bool,
+}
+
+enum SetupSnapshotResult {
+    WithSnapshot(SnapshotSetupInfo),
+    NoSnapshot,
+}
+
 impl<S: Signer<T> + Send + 'static, T: SignerEventTrait + 'static> SignerTest<SpawnedSigner<S, T>> {
     pub fn new(num_signers: usize, initial_balances: Vec<(StacksAddress, u64)>) -> Self {
         Self::new_with_config_modifications(
@@ -232,35 +242,42 @@ impl<S: Signer<T> + Send + 'static, T: SignerEventTrait + 'static> SignerTest<Sp
                 vec![pk]
             });
 
-        let mut snapshot_exists = false;
+        let snapshot_setup_result = Self::setup_snapshot(snapshot_name, &naka_conf);
 
-        let snapshot_path = snapshot_name.map(|name| {
-            let working_dir = naka_conf.get_working_dir();
+        // let mut snapshot_exists = false;
 
-            let snapshot_path: PathBuf = format!("/tmp/stacks-node-tests/snapshots/{name}/")
-                .try_into()
-                .unwrap();
+        // let snapshot_path = snapshot_name.map(|name| {
+        //     let working_dir = naka_conf.get_working_dir();
 
-            info!("Snapshot path: {}", snapshot_path.clone().display());
-
-            snapshot_exists = std::fs::metadata(snapshot_path.clone()).is_ok();
-
-            if snapshot_exists {
-                info!(
-                    "Snapshot directory already exists, copying to working dir";
-                    "snapshot_path" => %snapshot_path.display(),
-                    "working_dir" => %working_dir.display()
-                );
-                let err_msg = format!(
-                    "Failed to copy snapshot dir to working dir: {} -> {}",
-                    snapshot_path.display(),
-                    working_dir.display()
-                );
-                copy_dir_all(snapshot_path.clone(), working_dir).expect(&err_msg);
-            }
+        //     let snapshot_path: PathBuf = format!("/tmp/stacks-node-tests/snapshots/{name}/")
+        //         .try_into()
+        //         .unwrap();
 
-            snapshot_path
-        });
+        //     info!("Snapshot path: {}", snapshot_path.clone().display());
+
+        //     snapshot_exists = std::fs::metadata(snapshot_path.clone()).is_ok();
+
+        //     if snapshot_exists {
+        //         info!(
+        //             "Snapshot directory already exists, copying to working dir";
+        //             "snapshot_path" => %snapshot_path.display(),
+        //             "working_dir" => %working_dir.display()
+        //         );
+        //         let err_msg = format!(
+        //             "Failed to copy snapshot dir to working dir: {} -> {}",
+        //             snapshot_path.display(),
+        //             working_dir.display()
+        //         );
+        //         copy_dir_all(snapshot_path.clone(), working_dir).expect(&err_msg);
+        //     }
+
+        //     snapshot_path
+        // });
+
+        let snapshot_exists = match &snapshot_setup_result {
+            SetupSnapshotResult::WithSnapshot(info) => info.snapshot_exists,
+            SetupSnapshotResult::NoSnapshot => false,
+        };
 
         let node = setup_stx_btc_node(
             naka_conf,
@@ -280,7 +297,10 @@ impl<S: Signer<T> + Send + 'static, T: SignerEventTrait + 'static> SignerTest<Sp
             stacks_client,
             num_stacking_cycles: 12_u64,
             signer_configs,
-            snapshot_path,
+            snapshot_path: match &snapshot_setup_result {
+                SetupSnapshotResult::WithSnapshot(info) => Some(info.snapshot_path.clone()),
+                SetupSnapshotResult::NoSnapshot => None,
+            },
         }
     }
 
@@ -302,6 +322,48 @@ impl<S: Signer<T> + Send + 'static, T: SignerEventTrait + 'static> SignerTest<Sp
         !std::fs::metadata(snapshot_path).is_ok()
     }
 
+    /// Setup a snapshot by copying the snapshot directory to the working directory.
+    ///
+    /// If the env variable `STACKS_TEST_SNAPSHOT` is not set, this will return `NoSnapshot`.
+    fn setup_snapshot(snapshot_name: Option<&str>, conf: &NeonConfig) -> SetupSnapshotResult {
+        let Some(snapshot_name) = snapshot_name else {
+            return SetupSnapshotResult::NoSnapshot;
+        };
+
+        if env::var("STACKS_TEST_SNAPSHOT") != Ok("1".into()) {
+            return SetupSnapshotResult::NoSnapshot;
+        }
+
+        let working_dir = conf.get_working_dir();
+
+        let snapshot_path: PathBuf = format!("/tmp/stacks-node-tests/snapshots/{snapshot_name}/")
+            .try_into()
+            .unwrap();
+
+        info!("Snapshot path: {}", snapshot_path.clone().display());
+
+        let snapshot_exists = std::fs::metadata(snapshot_path.clone()).is_ok();
+
+        if snapshot_exists {
+            info!(
+                "Snapshot directory already exists, copying to working dir";
+                "snapshot_path" => %snapshot_path.display(),
+                "working_dir" => %working_dir.display()
+            );
+            let err_msg = format!(
+                "Failed to copy snapshot dir to working dir: {} -> {}",
+                snapshot_path.display(),
+                working_dir.display()
+            );
+            copy_dir_all(snapshot_path.clone(), working_dir).expect(&err_msg);
+        }
+
+        SetupSnapshotResult::WithSnapshot(SnapshotSetupInfo {
+            snapshot_path,
+            snapshot_exists,
+        })
+    }
+
     /// Make a snapshot of the current working directory.
     ///
     /// This will stop the bitcoind node and copy the working directory to the snapshot path.

testnet/stacks-node/src/tests/signer/v0.rs

@@ -324,8 +324,16 @@ impl SignerTest<SpawnedSigner> {
     /// This will also shutdown the test early, requiring a restart
     /// of the test.
     ///
+    /// If the test is not configured to use snapshots, it will boot to epoch 3.0
+    /// and continue.
+    ///
     /// Returns `true` if the snapshot was created.
     pub fn bootstrap_snapshot(&self) -> bool {
+        if self.snapshot_path.is_none() {
+            self.boot_to_epoch_3();
+            return false;
+        }
+
         if self.needs_snapshot() {
             self.boot_to_epoch_3();
             self.make_snapshot();