Network arg (#70)

* Allow passing network arg * Remove use_latest_elf * CitreaMode * Rename CitreaMode variant --------- Co-authored-by: jfldde <[email protected]>
chainwayxyz · Jan 10, 2025 · af85eae · af85eae
1 parent 0d1bee3
commit af85eae
Show file tree

Hide file tree

Showing 5 changed files with 35 additions and 6 deletions.
diff --git a/src/citrea_config/batch_prover.rs b/src/citrea_config/batch_prover.rs
@@ -39,8 +39,6 @@ pub struct BatchProverConfig {
     pub proof_sampling_number: usize,
     /// If true prover will try to recover ongoing proving sessions
     pub enable_recovery: bool,
-    /// Wether to always use the latest ELF or not
-    pub use_latest_elf: bool,
 }
 
 impl Default for BatchProverConfig {
@@ -49,7 +47,6 @@ impl Default for BatchProverConfig {
             proving_mode: ProverGuestRunConfig::Execute,
             proof_sampling_number: 0,
             enable_recovery: true,
-            use_latest_elf: true,
         }
     }
 }
@@ -100,7 +97,6 @@ mod tests {
             proving_mode: ProverGuestRunConfig::Skip,
             proof_sampling_number: 500,
             enable_recovery: true,
-            use_latest_elf: true,
         };
         assert_eq!(config, expected);
     }

diff --git a/src/config/mod.rs b/src/config/mod.rs
@@ -49,6 +49,23 @@ impl fmt::Display for DaLayer {
     }
 }
 
+#[derive(Clone, Debug, Default, Copy)]
+pub enum CitreaMode {
+    #[default]
+    Dev,
+    DevAllForks,
+}
+
+impl fmt::Display for CitreaMode {
+    // This trait requires `fmt` with this exact signature.
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            CitreaMode::Dev => write!(f, "dev"),
+            CitreaMode::DevAllForks => write!(f, "dev-all-forks"),
+        }
+    }
+}
+
 #[derive(Clone, Debug, Default)]
 pub struct FullL2NodeConfig<T> {
     pub node: T,
@@ -57,6 +74,7 @@ pub struct FullL2NodeConfig<T> {
     pub dir: PathBuf,
     pub env: Vec<(&'static str, &'static str)>,
     pub da_layer: Option<DaLayer>,
+    pub mode: CitreaMode,
 }
 
 impl<T> FullL2NodeConfig<T>
@@ -70,6 +88,7 @@ where
         docker_image: Option<String>,
         dir: PathBuf,
         env: Vec<(&'static str, &'static str)>,
+        mode: CitreaMode,
     ) -> Result<Self> {
         let conf = Self {
             node,
@@ -78,6 +97,7 @@ where
             dir,
             env,
             da_layer: None,
+            mode,
         };
 
         let kind = FullL2NodeConfig::<T>::kind();
@@ -190,6 +210,10 @@ where
     fn da_layer(&self) -> DaLayer {
         self.da_layer.clone().unwrap_or_default()
     }
+
+    fn mode(&self) -> CitreaMode {
+        self.mode
+    }
 }
 
 impl<T> LogPathProvider for FullL2NodeConfig<T>

diff --git a/src/config/test_case.rs b/src/config/test_case.rs
@@ -4,6 +4,8 @@ use tempfile::TempDir;
 
 use crate::utils::generate_test_id;
 
+use super::CitreaMode;
+
 #[derive(Clone, Default)]
 pub struct TestCaseEnv {
     pub test: Vec<(&'static str, &'static str)>,
@@ -60,6 +62,7 @@ pub struct TestCaseConfig {
     // Defaults to resources/genesis/bitcoin-regtest
     pub genesis_dir: Option<String>,
     pub test_id: String,
+    pub mode: CitreaMode,
 }
 
 impl Default for TestCaseConfig {
@@ -85,6 +88,7 @@ impl Default for TestCaseConfig {
             docker: TestCaseDockerConfig::default(),
             genesis_dir: None,
             test_id,
+            mode: CitreaMode::Dev,
         }
     }
 }

diff --git a/src/framework.rs b/src/framework.rs
@@ -443,27 +443,31 @@ fn generate_test_config<T: TestCase>(
             citrea_docker_image.clone(),
             sequencer_dir,
             env.sequencer(),
+            test_case.mode,
         )?,
         batch_prover: FullBatchProverConfig::new(
             batch_prover,
             batch_prover_rollup,
             citrea_docker_image.clone(),
             batch_prover_dir,
             env.batch_prover(),
+            test_case.mode,
         )?,
         light_client_prover: FullLightClientProverConfig::new(
             light_client_prover,
             light_client_prover_rollup,
             citrea_docker_image.clone(),
             light_client_prover_dir,
             env.light_client_prover(),
+            test_case.mode,
         )?,
         full_node: FullFullNodeConfig::new(
             (),
             full_node_rollup,
             citrea_docker_image,
             full_node_dir,
             env.full_node(),
+            test_case.mode,
         )?,
         test_case,
     })

diff --git a/src/node.rs b/src/node.rs
@@ -22,7 +22,7 @@ use tracing::{debug, info, trace};
 
 use crate::{
     client::Client,
-    config::{BitcoinConfig, DaLayer, DockerConfig, RollupConfig},
+    config::{BitcoinConfig, CitreaMode, DaLayer, DockerConfig, RollupConfig},
     docker::DockerEnv,
     log_provider::LogPathProvider,
     traits::{NodeT, Restart, SpawnOutput},
@@ -80,6 +80,7 @@ pub trait Config: Clone {
     // Not required for `full-node`
     fn get_node_config_args(&self) -> Option<Vec<String>>;
     fn get_rollup_config_args(&self) -> Vec<String>;
+    fn mode(&self) -> CitreaMode;
 }
 
 pub struct Node<C: Config + LogPathProvider + Send + Sync> {
@@ -284,7 +285,7 @@ where
     let rollup_config_args = config.get_rollup_config_args();
 
     [
-        vec!["--dev".to_string()],
+        vec![format!("--{}", config.mode())],
         vec!["--da-layer".to_string(), config.da_layer().to_string()],
         node_config_args,
         rollup_config_args,