Skip to main content

reth_e2e_test_utils/testsuite/
mod.rs

1//! Utilities for running e2e tests against a node or a network of nodes.
2
3use crate::{
4    testsuite::actions::{Action, ActionBox},
5    NodeBuilderHelper,
6};
7use alloy_primitives::{Bytes, B256};
8use eyre::Result;
9use jsonrpsee::http_client::HttpClient;
10use reth_node_api::{EngineTypes, PayloadTypes};
11use reth_payload_builder::{PayloadBuilderHandle, PayloadId};
12use std::{collections::HashMap, marker::PhantomData};
13pub mod actions;
14pub mod setup;
15use crate::testsuite::setup::Setup;
16use alloy_provider::{Provider, ProviderBuilder};
17use alloy_rpc_types_engine::{ForkchoiceState, PayloadAttributes};
18use reth_engine_primitives::ConsensusEngineHandle;
19use reth_rpc_builder::auth::AuthServerHandle;
20use std::sync::Arc;
21use url::Url;
22
23/// Client handles for both regular RPC and Engine API endpoints
24#[derive(Clone)]
25pub struct NodeClient<Payload>
26where
27    Payload: PayloadTypes,
28{
29    /// Regular JSON-RPC client
30    pub rpc: HttpClient,
31    /// Engine API client
32    pub engine: AuthServerHandle,
33    /// Beacon consensus engine handle for direct interaction with the consensus engine
34    pub beacon_engine_handle: Option<ConsensusEngineHandle<Payload>>,
35    /// Local payload builder used to wait for an in-progress build before requesting it over RPC.
36    pub(crate) payload_builder: Option<PayloadBuilderHandle<Payload>>,
37    /// Alloy provider for interacting with the node
38    provider: Arc<dyn Provider + Send + Sync>,
39}
40
41impl<Payload> NodeClient<Payload>
42where
43    Payload: PayloadTypes,
44{
45    /// Instantiates a new [`NodeClient`] with the given handles and RPC URL
46    pub fn new(rpc: HttpClient, engine: AuthServerHandle, url: Url) -> Self {
47        let provider =
48            Arc::new(ProviderBuilder::new().connect_http(url)) as Arc<dyn Provider + Send + Sync>;
49        Self { rpc, engine, beacon_engine_handle: None, payload_builder: None, provider }
50    }
51
52    /// Instantiates a new [`NodeClient`] with the given handles, RPC URL, and beacon engine handle
53    pub fn new_with_beacon_engine(
54        rpc: HttpClient,
55        engine: AuthServerHandle,
56        url: Url,
57        beacon_engine_handle: ConsensusEngineHandle<Payload>,
58    ) -> Self {
59        let provider =
60            Arc::new(ProviderBuilder::new().connect_http(url)) as Arc<dyn Provider + Send + Sync>;
61        Self {
62            rpc,
63            engine,
64            beacon_engine_handle: Some(beacon_engine_handle),
65            payload_builder: None,
66            provider,
67        }
68    }
69
70    /// Get a block by number using the alloy provider
71    pub async fn get_block_by_number(
72        &self,
73        number: alloy_eips::BlockNumberOrTag,
74    ) -> Result<Option<alloy_rpc_types_eth::Block>> {
75        self.provider
76            .get_block_by_number(number)
77            .await
78            .map_err(|e| eyre::eyre!("Failed to get block by number: {}", e))
79    }
80
81    /// Submit a raw transaction using the alloy provider.
82    pub async fn send_raw_transaction(&self, raw_tx: Bytes) -> Result<B256> {
83        let pending = self
84            .provider
85            .send_raw_transaction(&raw_tx)
86            .await
87            .map_err(|e| eyre::eyre!("Failed to send raw transaction: {}", e))?;
88        Ok(*pending.tx_hash())
89    }
90
91    /// Check if the node is ready by attempting to get the latest block
92    pub async fn is_ready(&self) -> bool {
93        self.get_block_by_number(alloy_eips::BlockNumberOrTag::Latest).await.is_ok()
94    }
95}
96
97impl<Payload> std::fmt::Debug for NodeClient<Payload>
98where
99    Payload: PayloadTypes,
100{
101    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
102        f.debug_struct("NodeClient")
103            .field("rpc", &self.rpc)
104            .field("engine", &self.engine)
105            .field("beacon_engine_handle", &self.beacon_engine_handle.is_some())
106            .field("provider", &"<Provider>")
107            .finish()
108    }
109}
110
111/// Represents complete block information.
112#[derive(Debug, Clone, Copy)]
113pub struct BlockInfo {
114    /// Hash of the block
115    pub hash: B256,
116    /// Number of the block
117    pub number: u64,
118    /// Timestamp of the block
119    pub timestamp: u64,
120}
121
122/// Per-node state tracking for multi-node environments
123#[derive(Clone)]
124pub struct NodeState<I>
125where
126    I: EngineTypes,
127{
128    /// Current block information for this node
129    pub current_block_info: Option<BlockInfo>,
130    /// Stores payload attributes indexed by block number for this node
131    pub payload_attributes: HashMap<u64, PayloadAttributes>,
132    /// Tracks the latest block header timestamp for this node
133    pub latest_header_time: u64,
134    /// Stores payload IDs returned by this node, indexed by block number
135    pub payload_id_history: HashMap<u64, PayloadId>,
136    /// Stores the next expected payload ID for this node
137    pub next_payload_id: Option<PayloadId>,
138    /// Stores the latest fork choice state for this node
139    pub latest_fork_choice_state: ForkchoiceState,
140    /// Stores the most recent built execution payload for this node
141    pub latest_payload_built: Option<PayloadAttributes>,
142    /// Stores the most recent executed payload for this node
143    pub latest_payload_executed: Option<PayloadAttributes>,
144    /// Stores the most recent built execution payload envelope for this node
145    pub latest_payload_envelope: Option<I::ExecutionPayloadEnvelopeV3>,
146    /// Fork base block number for validation (if this node is currently on a fork)
147    pub current_fork_base: Option<u64>,
148}
149
150impl<I> Default for NodeState<I>
151where
152    I: EngineTypes,
153{
154    fn default() -> Self {
155        Self {
156            current_block_info: None,
157            payload_attributes: HashMap::new(),
158            latest_header_time: 0,
159            payload_id_history: HashMap::new(),
160            next_payload_id: None,
161            latest_fork_choice_state: ForkchoiceState::default(),
162            latest_payload_built: None,
163            latest_payload_executed: None,
164            latest_payload_envelope: None,
165            current_fork_base: None,
166        }
167    }
168}
169
170impl<I> std::fmt::Debug for NodeState<I>
171where
172    I: EngineTypes,
173{
174    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
175        f.debug_struct("NodeState")
176            .field("current_block_info", &self.current_block_info)
177            .field("payload_attributes", &self.payload_attributes)
178            .field("latest_header_time", &self.latest_header_time)
179            .field("payload_id_history", &self.payload_id_history)
180            .field("next_payload_id", &self.next_payload_id)
181            .field("latest_fork_choice_state", &self.latest_fork_choice_state)
182            .field("latest_payload_built", &self.latest_payload_built)
183            .field("latest_payload_executed", &self.latest_payload_executed)
184            .field("latest_payload_envelope", &"<ExecutionPayloadEnvelopeV3>")
185            .field("current_fork_base", &self.current_fork_base)
186            .finish()
187    }
188}
189
190/// Represents a test environment.
191#[derive(Debug)]
192pub struct Environment<I>
193where
194    I: EngineTypes,
195{
196    /// Combined clients with both RPC and Engine API endpoints
197    pub node_clients: Vec<NodeClient<I>>,
198    /// Per-node state tracking
199    pub node_states: Vec<NodeState<I>>,
200    /// Tracks instance generic.
201    _phantom: PhantomData<I>,
202    /// Last producer index
203    pub last_producer_idx: Option<usize>,
204    /// Defines the increment for block timestamps (default: 2 seconds)
205    pub block_timestamp_increment: u64,
206    /// Number of slots until a block is considered safe
207    pub slots_to_safe: u64,
208    /// Number of slots until a block is considered finalized
209    pub slots_to_finalized: u64,
210    /// Registry for tagged blocks, mapping tag names to block info and node index
211    pub block_registry: HashMap<String, (BlockInfo, usize)>,
212    /// Currently active node index for backward compatibility with single-node actions
213    pub active_node_idx: usize,
214}
215
216impl<I> Default for Environment<I>
217where
218    I: EngineTypes,
219{
220    fn default() -> Self {
221        Self {
222            node_clients: vec![],
223            node_states: vec![],
224            _phantom: Default::default(),
225            last_producer_idx: None,
226            block_timestamp_increment: 2,
227            slots_to_safe: 0,
228            slots_to_finalized: 0,
229            block_registry: HashMap::new(),
230            active_node_idx: 0,
231        }
232    }
233}
234
235impl<I> Environment<I>
236where
237    I: EngineTypes,
238{
239    /// Get the number of nodes in the environment
240    pub const fn node_count(&self) -> usize {
241        self.node_clients.len()
242    }
243
244    /// Get mutable reference to a specific node's state
245    pub fn node_state_mut(&mut self, node_idx: usize) -> Result<&mut NodeState<I>, eyre::Error> {
246        let node_count = self.node_count();
247        self.node_states.get_mut(node_idx).ok_or_else(|| {
248            eyre::eyre!("Node index {} out of bounds (have {} nodes)", node_idx, node_count)
249        })
250    }
251
252    /// Get immutable reference to a specific node's state
253    pub fn node_state(&self, node_idx: usize) -> Result<&NodeState<I>, eyre::Error> {
254        self.node_states.get(node_idx).ok_or_else(|| {
255            eyre::eyre!("Node index {} out of bounds (have {} nodes)", node_idx, self.node_count())
256        })
257    }
258
259    /// Get the currently active node's state
260    pub fn active_node_state(&self) -> Result<&NodeState<I>, eyre::Error> {
261        self.node_state(self.active_node_idx)
262    }
263
264    /// Get mutable reference to the currently active node's state
265    pub fn active_node_state_mut(&mut self) -> Result<&mut NodeState<I>, eyre::Error> {
266        let idx = self.active_node_idx;
267        self.node_state_mut(idx)
268    }
269
270    /// Set the active node index
271    pub fn set_active_node(&mut self, node_idx: usize) -> Result<(), eyre::Error> {
272        if node_idx >= self.node_count() {
273            return Err(eyre::eyre!(
274                "Node index {} out of bounds (have {} nodes)",
275                node_idx,
276                self.node_count()
277            ));
278        }
279        self.active_node_idx = node_idx;
280        Ok(())
281    }
282
283    /// Initialize node states when nodes are created
284    pub fn initialize_node_states(&mut self, node_count: usize) {
285        self.node_states = (0..node_count).map(|_| NodeState::default()).collect();
286    }
287
288    /// Get current block info from active node
289    pub fn current_block_info(&self) -> Option<BlockInfo> {
290        self.active_node_state().ok()?.current_block_info
291    }
292
293    /// Set current block info on active node
294    pub fn set_current_block_info(&mut self, block_info: BlockInfo) -> Result<(), eyre::Error> {
295        self.active_node_state_mut()?.current_block_info = Some(block_info);
296        Ok(())
297    }
298}
299
300/// Builder for creating test scenarios
301#[expect(missing_debug_implementations)]
302pub struct TestBuilder<I>
303where
304    I: EngineTypes,
305{
306    setup: Option<Setup<I>>,
307    actions: Vec<ActionBox<I>>,
308    env: Environment<I>,
309}
310
311impl<I> Default for TestBuilder<I>
312where
313    I: EngineTypes,
314{
315    fn default() -> Self {
316        Self { setup: None, actions: Vec::new(), env: Default::default() }
317    }
318}
319
320impl<I> TestBuilder<I>
321where
322    I: EngineTypes + 'static,
323{
324    /// Create a new test builder
325    pub fn new() -> Self {
326        Self::default()
327    }
328
329    /// Set the test setup
330    pub fn with_setup(mut self, setup: Setup<I>) -> Self {
331        self.setup = Some(setup);
332        self
333    }
334
335    /// Set the test setup with chain import from RLP file
336    pub fn with_setup_and_import(
337        mut self,
338        mut setup: Setup<I>,
339        rlp_path: impl Into<std::path::PathBuf>,
340    ) -> Self {
341        setup.import_rlp_path = Some(rlp_path.into());
342        self.setup = Some(setup);
343        self
344    }
345
346    /// Add an action to the test
347    pub fn with_action<A>(mut self, action: A) -> Self
348    where
349        A: Action<I>,
350    {
351        self.actions.push(ActionBox::<I>::new(action));
352        self
353    }
354
355    /// Add multiple actions to the test
356    pub fn with_actions<II, A>(mut self, actions: II) -> Self
357    where
358        II: IntoIterator<Item = A>,
359        A: Action<I>,
360    {
361        self.actions.extend(actions.into_iter().map(ActionBox::new));
362        self
363    }
364
365    /// Run the test scenario
366    pub async fn run<N>(mut self) -> Result<()>
367    where
368        N: NodeBuilderHelper<Payload = I>,
369    {
370        let mut setup = self.setup.take();
371
372        if let Some(ref mut s) = setup {
373            s.apply::<N>(&mut self.env).await?;
374        }
375
376        let actions = std::mem::take(&mut self.actions);
377
378        for action in actions {
379            action.execute(&mut self.env).await?;
380        }
381
382        // explicitly drop the setup to shutdown the nodes
383        // after all actions have completed
384        drop(setup);
385
386        Ok(())
387    }
388}