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