From 7a08605d7089659acc7d7f866c6a857af08ef0aa Mon Sep 17 00:00:00 2001
From: jfldde <168934971+jfldde@users.noreply.github.com>
Date: Mon, 20 Jan 2025 10:44:22 +0000
Subject: [PATCH] Accept extra_args on node restart

---
 src/bitcoin.rs          |  8 ++++++--
 src/config/test_case.rs |  3 +--
 src/node.rs             | 13 +++++++++----
 src/traits.rs           | 14 +++++++++++---
 tests/bitcoin.rs        |  2 +-
 5 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/src/bitcoin.rs b/src/bitcoin.rs
index c2f1a81..c7daa85 100644
--- a/src/bitcoin.rs
+++ b/src/bitcoin.rs
@@ -7,7 +7,7 @@ use std::{
     time::{Duration, Instant},
 };
 
-use anyhow::{anyhow, bail, Context};
+use anyhow::{bail, Context};
 use async_trait::async_trait;
 use bitcoin::Address;
 use bitcoincore_rpc::{json::AddressType::Bech32m, Auth, Client, RpcApi};
@@ -283,7 +283,11 @@ impl Restart for BitcoinNode {
         }
     }
 
-    async fn start(&mut self, config: Option<Self::Config>) -> Result<()> {
+    async fn start(
+        &mut self,
+        config: Option<Self::Config>,
+        _extra_args: Option<Vec<String>>,
+    ) -> Result<()> {
         if let Some(config) = config {
             self.config = config;
         }
diff --git a/src/config/test_case.rs b/src/config/test_case.rs
index b4f09c5..dc8fa51 100644
--- a/src/config/test_case.rs
+++ b/src/config/test_case.rs
@@ -2,9 +2,8 @@ use std::{env, path::PathBuf, time::Duration};
 
 use tempfile::TempDir;
 
-use crate::utils::generate_test_id;
-
 use super::CitreaMode;
+use crate::utils::generate_test_id;
 
 #[derive(Clone, Default)]
 pub struct TestCaseEnv {
diff --git a/src/node.rs b/src/node.rs
index 3607f6a..1ac1acd 100644
--- a/src/node.rs
+++ b/src/node.rs
@@ -125,7 +125,7 @@ where
         })
     }
 
-    fn spawn(config: &C) -> Result<SpawnOutput> {
+    fn spawn(config: &C, extra_args: Option<Vec<String>>) -> Result<SpawnOutput> {
         let citrea = get_citrea_path()?;
 
         let kind = C::node_kind();
@@ -145,6 +145,7 @@ where
 
         Command::new(citrea)
             .args(get_citrea_args(config))
+            .args(extra_args.unwrap_or_default())
             .envs(config.env())
             .stdout(Stdio::from(stdout_file))
             .stderr(Stdio::from(stderr_file))
@@ -191,7 +192,7 @@ where
     async fn spawn(config: &Self::Config, docker: &Arc<Option<DockerEnv>>) -> Result<SpawnOutput> {
         match docker.as_ref() {
             Some(docker) if docker.citrea() => docker.spawn(config.to_owned().into()).await,
-            _ => Self::spawn(config),
+            _ => Self::spawn(config, None),
         }
     }
 
@@ -251,7 +252,11 @@ where
         Ok(())
     }
 
-    async fn start(&mut self, new_config: Option<Self::Config>) -> Result<()> {
+    async fn start(
+        &mut self,
+        new_config: Option<Self::Config>,
+        extra_args: Option<Vec<String>>,
+    ) -> Result<()> {
         let config = self.config_mut();
 
         if let Some(new_config) = new_config {
@@ -272,7 +277,7 @@ where
         copy_directory(old_dir, &new_dir)?;
         config.set_dir(new_dir);
 
-        *self.spawn_output() = Self::spawn(config)?;
+        *self.spawn_output() = Self::spawn(config, extra_args)?;
         self.wait_for_ready(None).await
     }
 }
diff --git a/src/traits.rs b/src/traits.rs
index 20c5b71..a1ff220 100644
--- a/src/traits.rs
+++ b/src/traits.rs
@@ -68,11 +68,19 @@ pub trait NodeT: Send {
 #[async_trait]
 pub trait Restart: NodeT + Send {
     async fn wait_until_stopped(&mut self) -> Result<()>;
-    async fn start(&mut self, new_config: Option<Self::Config>) -> Result<()>;
+    async fn start(
+        &mut self,
+        new_config: Option<Self::Config>,
+        extra_args: Option<Vec<String>>,
+    ) -> Result<()>;
 
     // Default implementation to support waiting for node to be fully shutdown and brough back up with new config.
-    async fn restart(&mut self, new_config: Option<Self::Config>) -> Result<()> {
+    async fn restart(
+        &mut self,
+        new_config: Option<Self::Config>,
+        extra_args: Option<Vec<String>>,
+    ) -> Result<()> {
         self.wait_until_stopped().await?;
-        self.start(new_config).await
+        self.start(new_config, extra_args).await
     }
 }
diff --git a/tests/bitcoin.rs b/tests/bitcoin.rs
index 424fbf6..23af334 100644
--- a/tests/bitcoin.rs
+++ b/tests/bitcoin.rs
@@ -108,7 +108,7 @@ impl TestCase for RestartBitcoinTest {
         assert_eq!(info.txindex, None);
 
         // Restart node with txindex
-        da.restart(Some(new_conf)).await?;
+        da.restart(Some(new_conf), None).await?;
 
         let block_after = da.get_block_count().await?;
         let info = da.get_index_info().await?;