feat(starknet_integration_tests): add db path argument to node_setup … (

#3925)

* chore(starknet_integration_tests): add support for specifying a custom database directory path

commit-id:675b21e5

* feat(starknet_integration_tests): add db path argument to node_setup for configurable DB storage

commit-id:2e49fc49

crates/papyrus_storage/src/test_utils.rs

@@ -1,6 +1,8 @@
 #![allow(clippy::unwrap_used)]
 //! Test utilities for the storage crate users.
 
+use std::fs::create_dir_all;
+use std::path::PathBuf;
 use std::sync::LazyLock;
 
 use starknet_api::core::ChainId;
@@ -14,30 +16,45 @@ use crate::{open_storage, StorageConfig, StorageReader, StorageScope, StorageWri
 pub static CHAIN_ID_FOR_TESTS: LazyLock<ChainId> =
     LazyLock::new(|| ChainId::Other("CHAIN_ID_SUBDIR".to_owned()));
 
+fn build_storage_config(storage_scope: StorageScope, path_prefix: PathBuf) -> StorageConfig {
+    StorageConfig {
+        db_config: DbConfig {
+            path_prefix,
+            chain_id: CHAIN_ID_FOR_TESTS.clone(),
+            enforce_file_exists: false,
+            min_size: 1 << 20,    // 1MB
+            max_size: 1 << 35,    // 32GB
+            growth_step: 1 << 26, // 64MB
+        },
+        scope: storage_scope,
+        mmap_file_config: get_mmap_file_test_config(),
+    }
+}
+
 /// Returns a db config and the temporary directory that holds this db.
 /// The TempDir object is returned as a handler for the lifetime of this object (the temp
 /// directory), thus make sure the directory won't be destroyed. The caller should propagate the
 /// TempDir object until it is no longer needed. When the TempDir object is dropped, the directory
 /// is deleted.
 pub(crate) fn get_test_config(storage_scope: Option<StorageScope>) -> (StorageConfig, TempDir) {
     let storage_scope = storage_scope.unwrap_or_default();
-    let dir = tempdir().unwrap();
-    println!("{dir:?}");
-    (
-        StorageConfig {
-            db_config: DbConfig {
-                path_prefix: dir.path().to_path_buf(),
-                chain_id: CHAIN_ID_FOR_TESTS.clone(),
-                enforce_file_exists: false,
-                min_size: 1 << 20,    // 1MB
-                max_size: 1 << 35,    // 32GB
-                growth_step: 1 << 26, // 64MB
-            },
-            scope: storage_scope,
-            mmap_file_config: get_mmap_file_test_config(),
-        },
-        dir,
-    )
+    let temp_dir = tempdir().expect("Failed to create temp directory");
+
+    (build_storage_config(storage_scope, temp_dir.path().to_path_buf()), temp_dir)
+}
+
+/// Returns a db config for a given path.
+/// This function ensures that the specified directory exists by creating it if necessary.
+/// Unlike get_test_config, this function does not return a TempDir object, so the caller is
+/// responsible for managing the directory's lifecycle and cleanup.
+pub(crate) fn get_test_config_with_path(
+    storage_scope: Option<StorageScope>,
+    path: PathBuf,
+) -> StorageConfig {
+    let storage_scope = storage_scope.unwrap_or_default();
+    create_dir_all(&path).expect("Failed to create directory");
+
+    build_storage_config(storage_scope, path)
 }
 
 /// Returns [`StorageReader`], [`StorageWriter`] and the temporary directory that holds a db for
@@ -84,7 +101,7 @@ pub fn get_test_storage_with_config_by_scope(
 /// A tool for creating a test storage.
 pub struct TestStorageBuilder {
     config: StorageConfig,
-    handle: TempDir,
+    handle: Option<TempDir>,
 }
 
 impl TestStorageBuilder {
@@ -104,15 +121,22 @@ impl TestStorageBuilder {
     /// that were built, and the temporary directory that holds a db for testing purposes. The
     /// returned [`StorageConfig`] can be used to open the exact same storage again (same DB
     /// file).
-    pub fn build(self) -> ((StorageReader, StorageWriter), StorageConfig, TempDir) {
+    pub fn build(self) -> ((StorageReader, StorageWriter), StorageConfig, Option<TempDir>) {
         let (reader, writer) = open_storage(self.config.clone()).unwrap();
         ((reader, writer), self.config, self.handle)
     }
-}
 
-impl Default for TestStorageBuilder {
-    fn default() -> Self {
-        let (config, handle) = get_test_config(None);
-        Self { config, handle }
+    /// Creates a TestStorageBuilder with a path to the directory that holds a db for testing.
+    pub fn new(path: Option<PathBuf>) -> Self {
+        match path {
+            Some(path) => {
+                let config = get_test_config_with_path(None, path);
+                Self { config, handle: None }
+            }
+            None => {
+                let (config, handle) = get_test_config(None);
+                Self { config, handle: Some(handle) }
+            }
+        }
     }
 }

crates/starknet_integration_tests/src/bin/sequencer_node_setup.rs

@@ -1,11 +1,19 @@
+use std::env::args;
+use std::path::PathBuf;
+
 use starknet_integration_tests::node_setup::node_setup;
 use starknet_integration_tests::utils::create_integration_test_tx_generator;
 use starknet_sequencer_infra::trace_util::configure_tracing;
 use tracing::info;
+
 #[tokio::main]
 async fn main() {
     configure_tracing().await;
-    info!("Running integration test setup.");
+    info!("Running system test setup.");
+
+    // Parse command line arguments.
+    let args: Vec<String> = args().skip(1).collect();
+    let base_db_path = get_base_db_path(args);
 
     // TODO(Tsabary): remove the hook definition once we transition to proper usage of task
     // spawning.
@@ -20,5 +28,19 @@ async fn main() {
 
     // Run node setup.
     // Keep the sequenser_setups in a variable to avoid dropping it.
-    let _sequencer_setups = node_setup(&mut tx_generator, "./single_node_config.json").await;
+    let _sequencer_setups =
+        node_setup(&mut tx_generator, "./single_node_config.json", base_db_path).await;
+}
+
+// TODO(Nadin): Improve the argument parsing.
+pub fn get_base_db_path(args: Vec<String>) -> PathBuf {
+    let arg_name = "--base_db_path_dir";
+    match args.as_slice() {
+        [] => PathBuf::from("./data"),
+        [arg, path] if arg == arg_name => PathBuf::from(path),
+        _ => {
+            eprintln!("Error: Bad argument. The only allowed argument is '{}'.", arg_name);
+            std::process::exit(1);
+        }
+    }
 }

crates/starknet_integration_tests/src/end_to_end_integration.rs

@@ -24,6 +24,7 @@ pub async fn end_to_end_integration(tx_generator: &mut MultiAccountTransactionGe
         tx_generator,
         N_CONSOLIDATED_SEQUENCERS,
         N_DISTRIBUTED_SEQUENCERS,
+        None,
     )
     .await;
 

crates/starknet_integration_tests/src/flow_test_setup.rs

@@ -122,8 +122,8 @@ pub struct FlowSequencerSetup {
     pub add_tx_http_client: HttpTestClient,
 
     // Handlers for the storage files, maintained so the files are not deleted.
-    pub batcher_storage_file_handle: TempDir,
-    pub state_sync_storage_file_handle: TempDir,
+    pub batcher_storage_file_handle: Option<TempDir>,
+    pub state_sync_storage_file_handle: Option<TempDir>,
 
     // Node configuration.
     pub node_config: SequencerNodeConfig,
@@ -153,7 +153,7 @@ impl FlowSequencerSetup {
             batcher_storage_handle,
             state_sync_storage_config,
             state_sync_storage_handle,
-        } = StorageTestSetup::new(accounts, &chain_info);
+        } = StorageTestSetup::new(accounts, &chain_info, None);
 
         let (recorder_url, _join_handle) =
             spawn_local_success_recorder(available_ports.get_next_port());

crates/starknet_integration_tests/src/integration_test_setup.rs

@@ -1,5 +1,5 @@
 use std::net::SocketAddr;
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 
 use blockifier::context::ChainInfo;
 use mempool_test_utils::starknet_api_test_utils::AccountTransactionGenerator;
@@ -39,6 +39,11 @@ impl NodeExecutionId {
     pub fn get_executable_index(&self) -> usize {
         self.executable_index
     }
+
+    pub fn build_path(&self, base: &Path) -> PathBuf {
+        base.join(format!("node_{}", self.node_index))
+            .join(format!("executable_{}", self.executable_index))
+    }
 }
 
 impl std::fmt::Display for NodeExecutionId {
@@ -70,11 +75,11 @@ pub struct ExecutableSetup {
     // these are only maintained to avoid dropping the handlers, private visibility suffices, and
     // as such, the '#[allow(dead_code)]' attributes are used to suppress the warning.
     #[allow(dead_code)]
-    batcher_storage_handle: TempDir,
+    batcher_storage_handle: Option<TempDir>,
     #[allow(dead_code)]
     node_config_dir_handle: TempDir,
     #[allow(dead_code)]
-    state_sync_storage_handle: TempDir,
+    state_sync_storage_handle: Option<TempDir>,
 }
 
 // TODO(Tsabary/ Nadin): reduce number of args.
@@ -90,6 +95,7 @@ impl ExecutableSetup {
         state_sync_config: StateSyncConfig,
         mut available_ports: AvailablePorts,
         component_config: ComponentConfig,
+        db_path_dir: Option<PathBuf>,
     ) -> Self {
         // TODO(Nadin): pass the test storage as an argument.
         // Creating the storage for the test.
@@ -98,7 +104,7 @@ impl ExecutableSetup {
             batcher_storage_handle,
             state_sync_storage_config,
             state_sync_storage_handle,
-        } = StorageTestSetup::new(accounts, &chain_info);
+        } = StorageTestSetup::new(accounts, &chain_info, db_path_dir);
 
         let (recorder_url, _join_handle) =
             spawn_local_success_recorder(available_ports.get_next_port());

crates/starknet_integration_tests/src/node_setup.rs

@@ -9,6 +9,7 @@ use crate::sequencer_manager::{get_sequencer_setup_configs, NodeSetup};
 pub async fn node_setup(
     tx_generator: &mut MultiAccountTransactionGenerator,
     config_path: &str,
+    base_db_path_dir: PathBuf,
 ) -> Vec<NodeSetup> {
     const N_CONSOLIDATED_SEQUENCERS: usize = 1;
     const N_DISTRIBUTED_SEQUENCERS: usize = 0;
@@ -19,6 +20,7 @@ pub async fn node_setup(
         tx_generator,
         N_CONSOLIDATED_SEQUENCERS,
         N_DISTRIBUTED_SEQUENCERS,
+        Some(base_db_path_dir),
     )
     .await;
 
@@ -31,3 +33,16 @@ pub async fn node_setup(
     println!("Config file moved from {:?} to {:?}", original_config_path, new_config_path);
     sequencers_setup
 }
+
+// TODO(Nadin): Improve the argument parsing.
+pub fn get_base_db_path(args: Vec<String>) -> PathBuf {
+    let arg_name = "--base_db_path_dir";
+    match args.as_slice() {
+        [] => PathBuf::from("./data"),
+        [arg, path] if arg == arg_name => PathBuf::from(path),
+        _ => {
+            eprintln!("Error: Bad argument. The only allowed argument is '{}'.", arg_name);
+            std::process::exit(1);
+        }
+    }
+}

crates/starknet_integration_tests/src/sequencer_manager.rs

@@ -1,5 +1,6 @@
 use std::collections::{HashMap, HashSet};
 use std::net::SocketAddr;
+use std::path::PathBuf;
 
 use futures::future::join_all;
 use futures::TryFutureExt;
@@ -310,6 +311,7 @@ pub(crate) async fn get_sequencer_setup_configs(
     tx_generator: &MultiAccountTransactionGenerator,
     num_of_consolidated_nodes: usize,
     num_of_distributed_nodes: usize,
+    path_to_base_dir: Option<PathBuf>,
 ) -> (Vec<NodeSetup>, HashSet<usize>) {
     let test_unique_id = TestIdentifier::EndToEndIntegrationTest;
 
@@ -376,6 +378,8 @@ pub(crate) async fn get_sequencer_setup_configs(
             let mempool_p2p_config = mempool_p2p_configs.remove(0);
             let state_sync_config = state_sync_configs.remove(0);
             let chain_info = chain_info.clone();
+            let exec_db_path = path_to_base_dir.as_ref().map(|p| node_execution_id.build_path(p));
+
             executables.push(
                 ExecutableSetup::new(
                     accounts.to_vec(),
@@ -386,6 +390,7 @@ pub(crate) async fn get_sequencer_setup_configs(
                     state_sync_config,
                     AvailablePorts::new(test_unique_id.into(), global_index.try_into().unwrap()),
                     executable_component_config.clone(),
+                    exec_db_path,
                 )
                 .await,
             );

crates/starknet_integration_tests/src/state_reader.rs

@@ -1,4 +1,5 @@
 use std::collections::HashMap;
+use std::path::PathBuf;
 
 use assert_matches::assert_matches;
 use blockifier::blockifier_versioned_constants::VersionedConstants;
@@ -46,27 +47,31 @@ type ContractClassesMap =
 
 pub struct StorageTestSetup {
     pub batcher_storage_config: StorageConfig,
-    pub batcher_storage_handle: TempDir,
+    pub batcher_storage_handle: Option<TempDir>,
     pub state_sync_storage_config: StorageConfig,
-    pub state_sync_storage_handle: TempDir,
+    pub state_sync_storage_handle: Option<TempDir>,
 }
 
 impl StorageTestSetup {
     pub fn new(
         test_defined_accounts: Vec<AccountTransactionGenerator>,
         chain_info: &ChainInfo,
+        path: Option<PathBuf>,
     ) -> Self {
+        let batcher_db_path = path.as_ref().map(|p| p.join("batcher"));
         let ((_, mut batcher_storage_writer), batcher_storage_config, batcher_storage_handle) =
-            TestStorageBuilder::default()
+            TestStorageBuilder::new(batcher_db_path)
                 .scope(StorageScope::StateOnly)
                 .chain_id(chain_info.chain_id.clone())
                 .build();
         create_test_state(&mut batcher_storage_writer, chain_info, test_defined_accounts.clone());
+
+        let state_sync_db_path = path.as_ref().map(|p| p.join("state_sync"));
         let (
             (_, mut state_sync_storage_writer),
             state_sync_storage_config,
             state_sync_storage_handle,
-        ) = TestStorageBuilder::default()
+        ) = TestStorageBuilder::new(state_sync_db_path)
             .scope(StorageScope::FullArchive)
             .chain_id(chain_info.chain_id.clone())
             .build();