Merge pull request #20 from ChainSafe/assethub_zombienet-forking-integration-tests

Assethub zombienet forking integration tests

Author: dimartiro

crates/anvil-polkadot/Cargo.toml

@@ -163,4 +163,5 @@ op-alloy-rpc-types.workspace = true
 
 [features]
 default = []
-asm-keccak = ["alloy-primitives/asm-keccak"]
+asm-keccak = ["alloy-primitives/asm-keccak"]
+forking-tests = []

crates/anvil-polkadot/src/substrate_node/service/client.rs

@@ -1,6 +1,5 @@
 use crate::{
     AnvilNodeConfig,
-    config::ForkChoice,
     substrate_node::{
         genesis::DevelopmentGenesisBlockBuilder,
         lazy_loading::{
@@ -25,7 +24,7 @@ use polkadot_sdk::{
     sp_blockchain,
     sp_core::storage::well_known_keys::CODE,
     sp_keystore::KeystorePtr,
-    sp_runtime::{generic::SignedBlock, traits::Header as HeaderT},
+    sp_runtime::generic::SignedBlock,
     sp_storage::StorageKey,
 };
 use std::{collections::HashMap, sync::Arc, time::Duration};
@@ -42,8 +41,7 @@ pub fn new_client(
 ) -> Result<(Arc<Client>, Arc<Backend>, KeystorePtr, TaskManager), sc_service::error::Error> {
     let fork_config: Option<(Arc<dyn RPCClient<Block>>, Block)> =
         if let Some(fork_url) = &anvil_config.eth_rpc_url {
-            let (rpc_client, checkpoint_block) =
-                setup_fork(anvil_config, config, fork_url, genesis_num)?;
+            let (rpc_client, checkpoint_block) = setup_fork(config, fork_url, genesis_num)?;
             Some((rpc_client, checkpoint_block))
         } else {
             None
@@ -124,49 +122,8 @@ pub fn new_client(
     Ok((Arc::new(client), backend, keystore_container.keystore(), task_manager))
 }
 
-/// Resolves the block number to fork from, handling both positive and negative block numbers.
-/// Negative numbers are subtracted from the latest block number.
-#[allow(dead_code)]
-fn resolve_fork_block_number(
-    rpc_client: &Rpc<Block>,
-    fork_choice: &ForkChoice,
-) -> Result<u32, sp_blockchain::Error> {
-    match fork_choice {
-        ForkChoice::Block(block_number) => {
-            if *block_number < 0 {
-                // Get the latest block from the chain header
-                let latest_header = rpc_client
-                    .header(None)
-                    .map_err(|e| {
-                        sp_blockchain::Error::Backend(format!("failed to fetch latest header: {e}"))
-                    })?
-                    .ok_or_else(|| {
-                        sp_blockchain::Error::Backend("latest header not found".into())
-                    })?;
-
-                let latest_number: u32 = *latest_header.number();
-
-                let offset: u32 = block_number.abs().try_into().map_err(|_| {
-                    sp_blockchain::Error::Backend(format!(
-                        "Block number offset too large: {block_number}"
-                    ))
-                })?;
-
-                Ok(latest_number.saturating_sub(offset))
-            } else {
-                (*block_number).try_into().map_err(|_| {
-                    sp_blockchain::Error::Backend(format!(
-                        "Invalid fork block number: {block_number}"
-                    ))
-                })
-            }
-        }
-    }
-}
-
 /// Fetches the checkpoint block and sets up the chain spec for forking
 fn setup_fork(
-    _anvil_config: &AnvilNodeConfig,
     config: &mut sc_service::Configuration,
     fork_url: &str,
     genesis_num: u64,
@@ -189,38 +146,6 @@ fn setup_fork(
         sp_blockchain::Error::Backend(format!("failed to fetch system_properties: {e}"))
     })?;
 
-    // TODO refactor this to account for fork choice
-    // Get block hash from fork_choice config
-    // If no fork_choice is specified, we need to fetch the latest block and use its hash
-    // for all subsequent requests to avoid inconsistencies if a new block is mined between calls.
-    // let block_hash: <Block as BlockT>::Hash = if let Some(fork_choice) =
-    // &anvil_config.fork_choice {    // let block_num = resolve_fork_block_number(&rpc_client,
-    // fork_choice)?;     let block_num = genesis_num as u32;
-    //     rpc_client
-    //         .block_hash(Some(block_num))
-    //         .map_err(|e| {
-    //             sp_blockchain::Error::Backend(format!(
-    //                 "failed to fetch block hash for block {block_num}: {e}"
-    //             ))
-    //         })?
-    //         .ok_or_else(|| {
-    //             sp_blockchain::Error::Backend(format!("block hash not found for block
-    // {block_num}"))         })?
-    // } else {
-    //     // No fork_choice specified, fetch the latest block header and use its hash
-    //     let latest_header = rpc_client
-    //         .header(None)
-    //         .map_err(|e| {
-    //             sp_blockchain::Error::Backend(format!(
-    //                 "failed to fetch latest header for fork: {e}"
-    //             ))
-    //         })?
-    //         .ok_or_else(|| {
-    //             sp_blockchain::Error::Backend("latest header not found for fork".into())
-    //         })?;
-    //     latest_header.hash()
-    // };
-
     let block_num = genesis_num as u32;
     let block_hash = rpc_client
         .block_hash(Some(block_num))

crates/anvil-polkadot/src/substrate_node/service/mod.rs

@@ -1,5 +1,6 @@
 use crate::{
     AnvilNodeConfig,
+    config::ForkChoice,
     substrate_node::{
         lazy_loading::backend::Backend as LazyLoadingBackend,
         mining_engine::{MiningEngine, MiningMode, run_mining_engine},
@@ -198,6 +199,7 @@ pub fn new(
     if let Some(ref fork_url) = anvil_config.eth_rpc_url {
         // TODO ws is for local host, wss for remote (aka prod)
         let http_url = fork_url.replacen("https://", "ws://", 1).replacen("http://", "ws://", 1);
+        let fork_choice = anvil_config.fork_choice.clone();
         let storage_map = std::thread::spawn(move || -> eyre::Result<Result<u64, ()>> {
             let rt = TokioRtBuilder::new_current_thread()
                 .enable_all()
@@ -217,7 +219,24 @@ pub fn new(
                     .await
                     .unwrap()
                     .unwrap();
-                Ok(Ok(finalized_head_header.number.into()))
+                let finalized_block_number: u64 = finalized_head_header.number.into();
+
+                // Apply fork_choice if specified
+                let target_block_number = match fork_choice {
+                    Some(ForkChoice::Block(block_num)) => {
+                        if block_num < 0 {
+                            // Negative offset from latest finalized block
+                            let offset = (-block_num) as u64;
+                            finalized_block_number.saturating_sub(offset)
+                        } else {
+                            // Specific block number
+                            block_num as u64
+                        }
+                    }
+                    None => finalized_block_number,
+                };
+
+                Ok(Ok(target_block_number))
             })
         })
         .join()