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

reth_e2e_test_utils/testsuite/
mod.rs
