From 7bf8d3e4b1f0c1d6f961d0c7decf8fba56df9066 Mon Sep 17 00:00:00 2001
From: vid277 <vidyoots@gmail.com>
Date: Wed, 25 Jun 2025 15:13:41 -0700
Subject: [PATCH] feature: ensure chain state is consistent across nodes

---
 tests/Cargo.toml                        |  2 +-
 tests/src/bin/integration-tests/main.rs | 49 +++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/tests/Cargo.toml b/tests/Cargo.toml
index 9aafd56d..3c3d7013 100644
--- a/tests/Cargo.toml
+++ b/tests/Cargo.toml
@@ -31,7 +31,7 @@ uuid.workspace = true
 tracing-subscriber.workspace = true
 bip39.workspace = true
 clap.workspace = true
-
+tonic.workspace = true
 abci = { path = "../crates/abci" }
 protocol = { path = "../crates/protocol" }
 node = { path = "../crates/node" }
diff --git a/tests/src/bin/integration-tests/main.rs b/tests/src/bin/integration-tests/main.rs
index eb5d1544..6bfa1e3e 100644
--- a/tests/src/bin/integration-tests/main.rs
+++ b/tests/src/bin/integration-tests/main.rs
@@ -10,6 +10,7 @@ use node::key_manager::generate_keys_from_mnemonic;
 use node::wallet::{TaprootWallet, Wallet};
 use oracle::esplora::EsploraOracle;
 use oracle::oracle::Oracle;
+use tonic::transport::Channel;
 use types::proto::node_proto::node_control_client::NodeControlClient;
 use types::proto::node_proto::{
     CheckBalanceRequest, ConfirmWithdrawalRequest, CreateDepositIntentRequest, GetChainInfoRequest,
@@ -658,6 +659,9 @@ async fn run_consensus_test(
         }
     }
 
+    println!("🔗 Verifying chain state consistency across nodes");
+    verify_chain_state_consistency(&mut clients).await?;
+
     // Check final balances to verify transaction execution
     println!("🔍 Verifying transaction execution across nodes");
     let mut execution_consistent = true;
@@ -764,3 +768,48 @@ where
         tokio::time::sleep(poll_interval).await;
     }
 }
+
+async fn verify_chain_state_consistency(
+    clients: &mut [NodeControlClient<Channel>],
+) -> Result<(), Box<dyn std::error::Error>> {
+    use types::proto::node_proto::GetChainInfoRequest;
+
+    let mut reference_height: Option<u64> = None;
+    let mut reference_hash: Option<String> = None;
+
+    for (idx, client) in clients.iter_mut().enumerate() {
+        let info = client
+            .get_chain_info(GetChainInfoRequest {})
+            .await?
+            .into_inner();
+
+        println!(
+            "    🌐 Node {}: height={} | hash={}",
+            idx + 1,
+            info.latest_height,
+            info.latest_block_hash
+        );
+
+        if let Some(h) = reference_height {
+            if let Some(hash) = &reference_hash {
+                if h != info.latest_height || hash != &info.latest_block_hash {
+                    return Err(format!(
+                        "Chain state mismatch on node {} (height/hash differ)",
+                        idx + 1
+                    )
+                    .into());
+                }
+            }
+        } else {
+            reference_height = Some(info.latest_height);
+            reference_hash = Some(info.latest_block_hash);
+        }
+    }
+
+    println!(
+        "    ✅ Chain state is consistent across all {} nodes",
+        clients.len()
+    );
+
+    Ok(())
+}