feat(e2e): add beacon consensus handle to NodeClient (#18632)

crates/e2e-test-utils/Cargo.toml

@@ -31,6 +31,7 @@ reth-tokio-util.workspace = true
 reth-stages-types.workspace = true
 reth-network-peers.workspace = true
 reth-engine-local.workspace = true
+reth-engine-primitives.workspace = true
 reth-tasks.workspace = true
 reth-node-ethereum.workspace = true
 reth-ethereum-primitives.workspace = true

crates/e2e-test-utils/src/node.rs

@@ -305,4 +305,20 @@ where
     pub fn auth_server_handle(&self) -> AuthServerHandle {
         self.inner.auth_server_handle().clone()
     }
+
+    /// Creates a [`crate::testsuite::NodeClient`] from this test context.
+    ///
+    /// This helper method extracts the necessary handles and creates a client
+    /// that can interact with both the regular RPC and Engine API endpoints.
+    /// It automatically includes the beacon engine handle for direct consensus engine interaction.
+    pub fn to_node_client(&self) -> eyre::Result<crate::testsuite::NodeClient<Payload>> {
+        let rpc = self
+            .rpc_client()
+            .ok_or_else(|| eyre::eyre!("Failed to create HTTP RPC client for node"))?;
+        let auth = self.auth_server_handle();
+        let url = self.rpc_url();
+        let beacon_handle = self.inner.add_ons_handle.beacon_engine_handle.clone();
+
+        Ok(crate::testsuite::NodeClient::new_with_beacon_engine(rpc, auth, url, beacon_handle))
+    }
 }

crates/e2e-test-utils/src/testsuite/mod.rs

@@ -16,27 +16,48 @@ pub mod setup;
 use crate::testsuite::setup::Setup;
 use alloy_provider::{Provider, ProviderBuilder};
 use alloy_rpc_types_engine::{ForkchoiceState, PayloadAttributes};
+use reth_engine_primitives::ConsensusEngineHandle;
 use reth_rpc_builder::auth::AuthServerHandle;
 use std::sync::Arc;
 use url::Url;
 
 /// Client handles for both regular RPC and Engine API endpoints
 #[derive(Clone)]
-pub struct NodeClient {
+pub struct NodeClient<Payload>
+where
+    Payload: PayloadTypes,
+{
     /// Regular JSON-RPC client
     pub rpc: HttpClient,
     /// Engine API client
     pub engine: AuthServerHandle,
+    /// Beacon consensus engine handle for direct interaction with the consensus engine
+    pub beacon_engine_handle: Option<ConsensusEngineHandle<Payload>>,
     /// Alloy provider for interacting with the node
     provider: Arc<dyn Provider + Send + Sync>,
 }
 
-impl NodeClient {
+impl<Payload> NodeClient<Payload>
+where
+    Payload: PayloadTypes,
+{
     /// Instantiates a new [`NodeClient`] with the given handles and RPC URL
     pub fn new(rpc: HttpClient, engine: AuthServerHandle, url: Url) -> Self {
         let provider =
             Arc::new(ProviderBuilder::new().connect_http(url)) as Arc<dyn Provider + Send + Sync>;
-        Self { rpc, engine, provider }
+        Self { rpc, engine, beacon_engine_handle: None, provider }
+    }
+
+    /// Instantiates a new [`NodeClient`] with the given handles, RPC URL, and beacon engine handle
+    pub fn new_with_beacon_engine(
+        rpc: HttpClient,
+        engine: AuthServerHandle,
+        url: Url,
+        beacon_engine_handle: ConsensusEngineHandle<Payload>,
+    ) -> Self {
+        let provider =
+            Arc::new(ProviderBuilder::new().connect_http(url)) as Arc<dyn Provider + Send + Sync>;
+        Self { rpc, engine, beacon_engine_handle: Some(beacon_engine_handle), provider }
     }
 
     /// Get a block by number using the alloy provider
@@ -56,11 +77,15 @@ impl NodeClient {
     }
 }
 
-impl std::fmt::Debug for NodeClient {
+impl<Payload> std::fmt::Debug for NodeClient<Payload>
+where
+    Payload: PayloadTypes,
+{
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("NodeClient")
             .field("rpc", &self.rpc)
             .field("engine", &self.engine)
+            .field("beacon_engine_handle", &self.beacon_engine_handle.is_some())
             .field("provider", &"<Provider>")
             .finish()
     }
@@ -152,7 +177,7 @@ where
     I: EngineTypes,
 {
     /// Combined clients with both RPC and Engine API endpoints
-    pub node_clients: Vec<NodeClient>,
+    pub node_clients: Vec<NodeClient<I>>,
     /// Per-node state tracking
     pub node_states: Vec<NodeState<I>>,
     /// Tracks instance generic.
@@ -323,7 +348,7 @@ where
     /// Run the test scenario
     pub async fn run<N>(mut self) -> Result<()>
     where
-        N: NodeBuilderHelper,
+        N: NodeBuilderHelper<Payload = I>,
         LocalPayloadAttributesBuilder<N::ChainSpec>: PayloadAttributesBuilder<
             <<N as NodeTypes>::Payload as PayloadTypes>::PayloadAttributes,
         >,

crates/e2e-test-utils/src/testsuite/setup.rs

@@ -142,7 +142,7 @@ where
         rlp_path: &Path,
     ) -> Result<()>
     where
-        N: NodeBuilderHelper,
+        N: NodeBuilderHelper<Payload = I>,
         LocalPayloadAttributesBuilder<N::ChainSpec>: PayloadAttributesBuilder<
             <<N as NodeTypes>::Payload as PayloadTypes>::PayloadAttributes,
         >,
@@ -158,7 +158,7 @@ where
         rlp_path: &Path,
     ) -> Result<()>
     where
-        N: NodeBuilderHelper,
+        N: NodeBuilderHelper<Payload = I>,
         LocalPayloadAttributesBuilder<N::ChainSpec>: PayloadAttributesBuilder<
             <<N as NodeTypes>::Payload as PayloadTypes>::PayloadAttributes,
         >,
@@ -175,6 +175,7 @@ where
                 .ok_or_else(|| eyre!("Failed to create HTTP RPC client for node"))?;
             let auth = node.auth_server_handle();
             let url = node.rpc_url();
+            // TODO: Pass beacon_engine_handle once import system supports generic types
             node_clients.push(crate::testsuite::NodeClient::new(rpc, auth, url));
         }
 
@@ -189,7 +190,7 @@ where
     /// Apply the setup to the environment
     pub async fn apply<N>(&mut self, env: &mut Environment<I>) -> Result<()>
     where
-        N: NodeBuilderHelper,
+        N: NodeBuilderHelper<Payload = I>,
         LocalPayloadAttributesBuilder<N::ChainSpec>: PayloadAttributesBuilder<
             <<N as NodeTypes>::Payload as PayloadTypes>::PayloadAttributes,
         >,
@@ -201,7 +202,7 @@ where
     /// Apply the setup to the environment
     async fn apply_<N>(&mut self, env: &mut Environment<I>) -> Result<()>
     where
-        N: NodeBuilderHelper,
+        N: NodeBuilderHelper<Payload = I>,
         LocalPayloadAttributesBuilder<N::ChainSpec>: PayloadAttributesBuilder<
             <<N as NodeTypes>::Payload as PayloadTypes>::PayloadAttributes,
         >,
@@ -236,13 +237,7 @@ where
             Ok((nodes, executor, _wallet)) => {
                 // create HTTP clients for each node's RPC and Engine API endpoints
                 for node in &nodes {
-                    let rpc = node
-                        .rpc_client()
-                        .ok_or_else(|| eyre!("Failed to create HTTP RPC client for node"))?;
-                    let auth = node.auth_server_handle();
-                    let url = node.rpc_url();
-
-                    node_clients.push(crate::testsuite::NodeClient::new(rpc, auth, url));
+                    node_clients.push(node.to_node_client()?);
                 }
 
                 // spawn a separate task just to handle the shutdown
@@ -274,7 +269,7 @@ where
         rlp_path: &Path,
     ) -> Result<crate::setup_import::ChainImportResult>
     where
-        N: NodeBuilderHelper,
+        N: NodeBuilderHelper<Payload = I>,
         LocalPayloadAttributesBuilder<N::ChainSpec>: PayloadAttributesBuilder<
             <<N as NodeTypes>::Payload as PayloadTypes>::PayloadAttributes,
         >,
@@ -309,7 +304,7 @@ where
         &self,
     ) -> impl Fn(u64) -> <<N as NodeTypes>::Payload as PayloadTypes>::PayloadBuilderAttributes + Copy
     where
-        N: NodeBuilderHelper,
+        N: NodeBuilderHelper<Payload = I>,
         LocalPayloadAttributesBuilder<N::ChainSpec>: PayloadAttributesBuilder<
             <<N as NodeTypes>::Payload as PayloadTypes>::PayloadAttributes,
         >,
@@ -332,7 +327,7 @@ where
     async fn finalize_setup(
         &self,
         env: &mut Environment<I>,
-        node_clients: Vec<crate::testsuite::NodeClient>,
+        node_clients: Vec<crate::testsuite::NodeClient<I>>,
         use_latest_block: bool,
     ) -> Result<()> {
         if node_clients.is_empty() {
@@ -394,10 +389,13 @@ where
     }
 
     /// Wait for all nodes to be ready to accept RPC requests
-    async fn wait_for_nodes_ready(
+    async fn wait_for_nodes_ready<P>(
         &self,
-        node_clients: &[crate::testsuite::NodeClient],
-    ) -> Result<()> {
+        node_clients: &[crate::testsuite::NodeClient<P>],
+    ) -> Result<()>
+    where
+        P: PayloadTypes,
+    {
         for (idx, client) in node_clients.iter().enumerate() {
             let mut retry_count = 0;
             const MAX_RETRIES: usize = 10;
@@ -423,11 +421,14 @@ where
     }
 
     /// Get block info for a given block number or tag
-    async fn get_block_info(
+    async fn get_block_info<P>(
         &self,
-        client: &crate::testsuite::NodeClient,
+        client: &crate::testsuite::NodeClient<P>,
         block: BlockNumberOrTag,
-    ) -> Result<crate::testsuite::BlockInfo> {
+    ) -> Result<crate::testsuite::BlockInfo>
+    where
+        P: PayloadTypes,
+    {
         let block = client
             .get_block_by_number(block)
             .await?