diff --git a/crates/e2e-test-utils/src/testsuite/actions.rs b/crates/e2e-test-utils/src/testsuite/actions.rs index c459972a1a..9471632573 100644 --- a/crates/e2e-test-utils/src/testsuite/actions.rs +++ b/crates/e2e-test-utils/src/testsuite/actions.rs @@ -153,7 +153,75 @@ where }) } } +/// Pick the next block producer based on the latest block information. +#[derive(Debug)] +pub struct PickNextBlockProducer { + /// Tracks engine type + _phantom: PhantomData, +} +impl Default for PickNextBlockProducer { + fn default() -> Self { + Self::new() + } +} + +impl PickNextBlockProducer { + /// Create a new `PickNextBlockProducer` action + pub fn new() -> Self { + Self { _phantom: Default::default() } + } +} + +impl Action for PickNextBlockProducer +where + Engine: EngineTypes, +{ + fn execute<'a>(&'a mut self, env: &'a mut Environment) -> BoxFuture<'a, Result<()>> { + Box::pin(async move { + let num_clients = env.node_clients.len(); + if num_clients == 0 { + return Err(eyre::eyre!("No node clients available")); + } + + let latest_info = env + .latest_block_info + .as_ref() + .ok_or_else(|| eyre::eyre!("No latest block information available"))?; + + // Calculate the starting index based on the latest block number + let start_idx = ((latest_info.number + 1) % num_clients as u64) as usize; + + for i in 0..num_clients { + let idx = (start_idx + i) % num_clients; + let node_client = &env.node_clients[idx]; + let rpc_client = &node_client.rpc; + + let latest_block = + EthApiClient::::block_by_number( + rpc_client, + alloy_eips::BlockNumberOrTag::Latest, + false, + ) + .await?; + + if let Some(block) = latest_block { + let block_number = block.header.number; + let block_hash = block.header.hash; + + // Check if the block hash and number match the latest block info + if block_hash == latest_info.hash && block_number == latest_info.number { + env.last_producer_idx = Some(idx); + debug!("Selected node {} as the next block producer", idx); + return Ok(()); + } + } + } + + Err(eyre::eyre!("No suitable block producer found")) + }) + } +} /// Run a sequence of actions in series. #[allow(missing_debug_implementations)] pub struct Sequence { diff --git a/crates/e2e-test-utils/src/testsuite/mod.rs b/crates/e2e-test-utils/src/testsuite/mod.rs index 571588c267..e70fa372ae 100644 --- a/crates/e2e-test-utils/src/testsuite/mod.rs +++ b/crates/e2e-test-utils/src/testsuite/mod.rs @@ -4,6 +4,7 @@ use crate::{ testsuite::actions::{Action, ActionBox}, NodeBuilderHelper, PayloadAttributesBuilder, }; +use alloy_primitives::B256; use eyre::Result; use jsonrpsee::http_client::{transport::HttpBackend, HttpClient}; use reth_engine_local::LocalPayloadAttributesBuilder; @@ -11,7 +12,6 @@ use reth_node_api::{NodeTypesWithEngine, PayloadTypes}; use reth_rpc_layer::AuthClientService; use setup::Setup; use std::marker::PhantomData; - pub mod actions; pub mod setup; @@ -27,6 +27,14 @@ pub struct NodeClient { pub engine: HttpClient>, } +/// Represents the latest block information. +#[derive(Debug, Clone)] +pub struct LatestBlockInfo { + /// Hash of the latest block + pub hash: B256, + /// Number of the latest block + pub number: u64, +} /// Represents a test environment. #[derive(Debug)] pub struct Environment { @@ -34,11 +42,20 @@ pub struct Environment { pub node_clients: Vec, /// Tracks instance generic. _phantom: PhantomData, + /// Latest block information + pub latest_block_info: Option, + /// Last producer index + pub last_producer_idx: Option, } impl Default for Environment { fn default() -> Self { - Self { node_clients: vec![], _phantom: Default::default() } + Self { + node_clients: vec![], + _phantom: Default::default(), + latest_block_info: None, + last_producer_idx: None, + } } }