From a4f5c4c8af8c07cbbcaf325cd9685e0c7b414b72 Mon Sep 17 00:00:00 2001
From: Omar Abdulla <OmarAbdulla7@hotmail.com>
Date: Mon, 14 Jul 2025 22:16:37 +0300
Subject: [PATCH] Add ability to get the chain_id from node

---
 crates/node-interaction/src/lib.rs |  5 ++++-
 crates/node/src/geth.rs            | 33 ++++++++++++++++++++++++++++-
 crates/node/src/kitchensink.rs     | 34 +++++++++++++++++++++++++++++-
 3 files changed, 69 insertions(+), 3 deletions(-)

diff --git a/crates/node-interaction/src/lib.rs b/crates/node-interaction/src/lib.rs
index 2c2eef5..d73d838 100644
--- a/crates/node-interaction/src/lib.rs
+++ b/crates/node-interaction/src/lib.rs
@@ -1,6 +1,6 @@
 //! This crate implements all node interactions.
 
-use alloy::primitives::Address;
+use alloy::primitives::{Address, ChainId};
 use alloy::rpc::types::trace::geth::{DiffMode, GethTrace};
 use alloy::rpc::types::{TransactionReceipt, TransactionRequest};
 
@@ -23,4 +23,7 @@ pub trait EthereumNode {
 
     /// Returns the next available nonce for the given [Address].
     fn fetch_add_nonce(&self, address: Address) -> anyhow::Result<u64>;
+
+    /// Returns the ID of the chain that the node is on.
+    fn chain_id(&self) -> anyhow::Result<ChainId>;
 }
diff --git a/crates/node/src/geth.rs b/crates/node/src/geth.rs
index 52a8aed..e6f500a 100644
--- a/crates/node/src/geth.rs
+++ b/crates/node/src/geth.rs
@@ -352,6 +352,14 @@ impl EthereumNode for Instance {
         *current += 1;
         Ok(value)
     }
+
+    #[tracing::instrument(skip_all, fields(geth_node_id = self.id))]
+    fn chain_id(&self) -> anyhow::Result<alloy::primitives::ChainId> {
+        let provider = self.provider();
+        BlockingExecutor::execute(async move {
+            provider.await?.get_chain_id().await.map_err(Into::into)
+        })?
+    }
 }
 
 impl Node for Instance {
@@ -439,7 +447,7 @@ mod tests {
 
     use crate::{GENESIS_JSON, Node};
 
-    use super::Instance;
+    use super::*;
 
     fn test_config() -> (Arguments, TempDir) {
         let mut config = Arguments::default();
@@ -449,6 +457,16 @@ mod tests {
         (config, temp_dir)
     }
 
+    fn new_node() -> (Instance, TempDir) {
+        let (args, temp_dir) = test_config();
+        let mut node = Instance::new(&args);
+        node.init(GENESIS_JSON.to_owned())
+            .expect("Failed to initialize the node")
+            .spawn_process()
+            .expect("Failed to spawn the node process");
+        (node, temp_dir)
+    }
+
     #[test]
     fn init_works() {
         Instance::new(&test_config().0)
@@ -471,4 +489,17 @@ mod tests {
             "expected version string, got: '{version}'"
         );
     }
+
+    #[test]
+    fn can_get_chain_id_from_node() {
+        // Arrange
+        let (node, _temp_dir) = new_node();
+
+        // Act
+        let chain_id = node.chain_id();
+
+        // Assert
+        let chain_id = chain_id.expect("Failed to get the chain id");
+        assert_eq!(chain_id, 420_420_420);
+    }
 }
diff --git a/crates/node/src/kitchensink.rs b/crates/node/src/kitchensink.rs
index b8586a9..9a186fb 100644
--- a/crates/node/src/kitchensink.rs
+++ b/crates/node/src/kitchensink.rs
@@ -229,6 +229,7 @@ impl KitchensinkNode {
 
         Ok(())
     }
+
     #[tracing::instrument(skip_all, fields(kitchensink_node_id = self.id))]
     fn extract_balance_from_genesis_file(
         &self,
@@ -415,6 +416,14 @@ impl EthereumNode for KitchensinkNode {
         *current += 1;
         Ok(value)
     }
+
+    #[tracing::instrument(skip_all, fields(geth_node_id = self.id))]
+    fn chain_id(&self) -> anyhow::Result<alloy::primitives::ChainId> {
+        let provider = self.provider();
+        BlockingExecutor::execute(async move {
+            provider.await?.get_chain_id().await.map_err(Into::into)
+        })?
+    }
 }
 
 impl Node for KitchensinkNode {
@@ -507,7 +516,7 @@ mod tests {
 
     use std::fs;
 
-    use super::KitchensinkNode;
+    use super::*;
     use crate::{GENESIS_JSON, Node};
 
     fn test_config() -> (Arguments, TempDir) {
@@ -522,6 +531,16 @@ mod tests {
         (config, temp_dir)
     }
 
+    fn new_node() -> (KitchensinkNode, TempDir) {
+        let (args, temp_dir) = test_config();
+        let mut node = KitchensinkNode::new(&args);
+        node.init(GENESIS_JSON)
+            .expect("Failed to initialize the node")
+            .spawn_process()
+            .expect("Failed to spawn the node process");
+        (node, temp_dir)
+    }
+
     #[test]
     fn test_init_generates_chainspec_with_balances() {
         let genesis_content = r#"
@@ -683,4 +702,17 @@ mod tests {
             "Expected eth-rpc version string, got: {version}"
         );
     }
+
+    #[test]
+    fn can_get_chain_id_from_node() {
+        // Arrange
+        let (node, _temp_dir) = new_node();
+
+        // Act
+        let chain_id = node.chain_id();
+
+        // Assert
+        let chain_id = chain_id.expect("Failed to get the chain id");
+        assert_eq!(chain_id, 420_420_420);
+    }
 }