reth_e2e_test_utils/testsuite/
mod.rs1use 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#[derive(Clone)]
25pub struct NodeClient<Payload>
26where
27 Payload: PayloadTypes,
28{
29 pub rpc: HttpClient,
31 pub engine: AuthServerHandle,
33 pub beacon_engine_handle: Option<ConsensusEngineHandle<Payload>>,
35 pub(crate) payload_builder: Option<PayloadBuilderHandle<Payload>>,
37 provider: Arc<dyn Provider + Send + Sync>,
39}
40
41impl<Payload> NodeClient<Payload>
42where
43 Payload: PayloadTypes,
44{
45 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 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 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 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 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#[derive(Debug, Clone, Copy)]
113pub struct BlockInfo {
114 pub hash: B256,
116 pub number: u64,
118 pub timestamp: u64,
120}
121
122#[derive(Clone)]
124pub struct NodeState<I>
125where
126 I: EngineTypes,
127{
128 pub current_block_info: Option<BlockInfo>,
130 pub payload_attributes: HashMap<u64, PayloadAttributes>,
132 pub latest_header_time: u64,
134 pub payload_id_history: HashMap<u64, PayloadId>,
136 pub next_payload_id: Option<PayloadId>,
138 pub latest_fork_choice_state: ForkchoiceState,
140 pub latest_payload_built: Option<PayloadAttributes>,
142 pub latest_payload_executed: Option<PayloadAttributes>,
144 pub latest_payload_envelope: Option<I::ExecutionPayloadEnvelopeV3>,
146 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#[derive(Debug)]
192pub struct Environment<I>
193where
194 I: EngineTypes,
195{
196 pub node_clients: Vec<NodeClient<I>>,
198 pub node_states: Vec<NodeState<I>>,
200 _phantom: PhantomData<I>,
202 pub last_producer_idx: Option<usize>,
204 pub block_timestamp_increment: u64,
206 pub slots_to_safe: u64,
208 pub slots_to_finalized: u64,
210 pub block_registry: HashMap<String, (BlockInfo, usize)>,
212 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 pub const fn node_count(&self) -> usize {
241 self.node_clients.len()
242 }
243
244 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 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 pub fn active_node_state(&self) -> Result<&NodeState<I>, eyre::Error> {
261 self.node_state(self.active_node_idx)
262 }
263
264 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 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 pub fn initialize_node_states(&mut self, node_count: usize) {
285 self.node_states = (0..node_count).map(|_| NodeState::default()).collect();
286 }
287
288 pub fn current_block_info(&self) -> Option<BlockInfo> {
290 self.active_node_state().ok()?.current_block_info
291 }
292
293 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#[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 pub fn new() -> Self {
326 Self::default()
327 }
328
329 pub fn with_setup(mut self, setup: Setup<I>) -> Self {
331 self.setup = Some(setup);
332 self
333 }
334
335 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 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 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 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 drop(setup);
385
386 Ok(())
387 }
388}