Reorg support (#42)

src/bitcoin.rs

@@ -306,8 +306,9 @@ impl BitcoinNodeCluster {
         Ok(())
     }
 
-    pub async fn wait_for_sync(&self, timeout: Duration) -> Result<()> {
+    pub async fn wait_for_sync(&self, timeout: Option<Duration>) -> Result<()> {
         let start = Instant::now();
+        let timeout = timeout.unwrap_or(Duration::from_secs(30));
         while start.elapsed() < timeout {
             let mut heights = HashSet::new();
             for node in &self.inner {
@@ -326,16 +327,34 @@ impl BitcoinNodeCluster {
 
     // Connect all bitcoin nodes between them
     pub async fn connect_nodes(&self) -> Result<()> {
-        for (i, from_node) in self.inner.iter().enumerate() {
-            for (j, to_node) in self.inner.iter().enumerate() {
+        for (i, from_node) in self.iter().enumerate() {
+            for (j, to_node) in self.iter().enumerate() {
+                if i != j {
+                    let ip = match &to_node.spawn_output {
+                        SpawnOutput::Container(container) => container.ip.clone(),
+                        SpawnOutput::Child(_) => "127.0.0.1".to_string(),
+                    };
+
+                    let target_node_addr = format!("{}:{}", ip, to_node.config.p2p_port);
+                    from_node.onetry_node(&target_node_addr).await?;
+                }
+            }
+        }
+        Ok(())
+    }
+
+    pub async fn disconnect_nodes(&self) -> Result<()> {
+        for (i, from_node) in self.iter().enumerate() {
+            for (j, to_node) in self.iter().enumerate() {
                 if i != j {
                     let ip = match &to_node.spawn_output {
                         SpawnOutput::Container(container) => container.ip.clone(),
                         SpawnOutput::Child(_) => "127.0.0.1".to_string(),
                     };
 
-                    let add_node_arg = format!("{}:{}", ip, to_node.config.p2p_port);
-                    from_node.add_node(&add_node_arg).await?;
+                    let target_node_addr = format!("{}:{}", ip, to_node.config.p2p_port);
+                    // from_node.remove_node(&target_node_addr).await?;
+                    from_node.disconnect_node(&target_node_addr).await?;
                 }
             }
         }

src/config/docker.rs

@@ -1,13 +1,14 @@
-use std::fmt::Debug;
-use std::path::PathBuf;
+use std::{fmt::Debug, path::PathBuf};
 
 use serde::Serialize;
 use tracing::debug;
 
 use super::{BitcoinConfig, FullL2NodeConfig, NodeKindMarker};
-use crate::log_provider::LogPathProvider;
-use crate::node::{get_citrea_args, Config, NodeKind};
-use crate::utils::get_genesis_path;
+use crate::{
+    log_provider::LogPathProvider,
+    node::{get_citrea_args, Config, NodeKind},
+    utils::get_genesis_path,
+};
 
 const DEFAULT_BITCOIN_DOCKER_IMAGE: &str = "bitcoin/bitcoin:28.0";
 const DEFAULT_CITREA_DOCKER_IMAGE: &str =

src/docker.rs

@@ -18,9 +18,8 @@ use futures::StreamExt;
 use tokio::{fs::File, io::AsyncWriteExt, sync::Mutex, task::JoinHandle};
 use tracing::{debug, error, info};
 
-use crate::{config::TestCaseDockerConfig, node::NodeKind};
-
 use super::{config::DockerConfig, traits::SpawnOutput, utils::generate_test_id};
+use crate::{config::TestCaseDockerConfig, node::NodeKind};
 
 #[derive(Debug)]
 pub struct ContainerSpawnOutput {