From ad85025035019eda9fec4688df90866a75f17ca9 Mon Sep 17 00:00:00 2001
From: jfldde <168934971+jfldde@users.noreply.github.com>
Date: Thu, 9 Jan 2025 17:41:48 +0000
Subject: [PATCH] CitreaMode

---
 src/config/mod.rs       | 24 ++++++++++++++++++++++++
 src/config/test_case.rs |  4 ++++
 src/framework.rs        |  4 ++++
 src/node.rs             | 10 +++-------
 4 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/src/config/mod.rs b/src/config/mod.rs
index 264788c..25ee454 100644
--- 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,
+    AllForks,
+}
+
+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::AllForks => write!(f, "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
index 8e1e921..b4f09c5 100644
--- 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
index be315bf..c3dc94f 100644
--- a/src/framework.rs
+++ b/src/framework.rs
@@ -443,6 +443,7 @@ fn generate_test_config<T: TestCase>(
             citrea_docker_image.clone(),
             sequencer_dir,
             env.sequencer(),
+            test_case.mode,
         )?,
         batch_prover: FullBatchProverConfig::new(
             batch_prover,
@@ -450,6 +451,7 @@ fn generate_test_config<T: TestCase>(
             citrea_docker_image.clone(),
             batch_prover_dir,
             env.batch_prover(),
+            test_case.mode,
         )?,
         light_client_prover: FullLightClientProverConfig::new(
             light_client_prover,
@@ -457,6 +459,7 @@ fn generate_test_config<T: TestCase>(
             citrea_docker_image.clone(),
             light_client_prover_dir,
             env.light_client_prover(),
+            test_case.mode,
         )?,
         full_node: FullFullNodeConfig::new(
             (),
@@ -464,6 +467,7 @@ fn generate_test_config<T: TestCase>(
             citrea_docker_image,
             full_node_dir,
             env.full_node(),
+            test_case.mode,
         )?,
         test_case,
     })
diff --git a/src/node.rs b/src/node.rs
index 2ef4a80..3607f6a 100644
--- 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> {
@@ -283,13 +284,8 @@ where
     let node_config_args = config.get_node_config_args().unwrap_or_default();
     let rollup_config_args = config.get_rollup_config_args();
 
-    let network_arg = match std::env::var("CITREA_NETWORK") {
-        Ok(network) => vec!["--network".to_string(), network],
-        _ => vec!["--dev".to_string()],
-    };
-
     [
-        network_arg,
+        vec![format!("--{}", config.mode())],
         vec!["--da-layer".to_string(), config.da_layer().to_string()],
         node_config_args,
         rollup_config_args,