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

reth_e2e_test_utils/testsuite/
mod.rs
