reth_e2e_test_utils/testsuite/
mod.rs1use 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#[derive(Clone)]
25pub struct NodeClient {
26 pub rpc: HttpClient,
28 pub engine: AuthServerHandle,
30 provider: Arc<dyn Provider + Send + Sync>,
32}
33
34impl NodeClient {
35 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 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 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#[derive(Debug, Clone, Copy)]
71pub struct BlockInfo {
72 pub hash: B256,
74 pub number: u64,
76 pub timestamp: u64,
78}
79
80#[derive(Clone)]
82pub struct NodeState<I>
83where
84 I: EngineTypes,
85{
86 pub current_block_info: Option<BlockInfo>,
88 pub payload_attributes: HashMap<u64, PayloadAttributes>,
90 pub latest_header_time: u64,
92 pub payload_id_history: HashMap<u64, PayloadId>,
94 pub next_payload_id: Option<PayloadId>,
96 pub latest_fork_choice_state: ForkchoiceState,
98 pub latest_payload_built: Option<PayloadAttributes>,
100 pub latest_payload_executed: Option<PayloadAttributes>,
102 pub latest_payload_envelope: Option<I::ExecutionPayloadEnvelopeV3>,
104 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#[derive(Debug)]
150pub struct Environment<I>
151where
152 I: EngineTypes,
153{
154 pub node_clients: Vec<NodeClient>,
156 pub node_states: Vec<NodeState<I>>,
158 _phantom: PhantomData<I>,
160 pub last_producer_idx: Option<usize>,
162 pub block_timestamp_increment: u64,
164 pub slots_to_safe: u64,
166 pub slots_to_finalized: u64,
168 pub block_registry: HashMap<String, (BlockInfo, usize)>,
170 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 pub const fn node_count(&self) -> usize {
199 self.node_clients.len()
200 }
201
202 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 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 pub fn active_node_state(&self) -> Result<&NodeState<I>, eyre::Error> {
219 self.node_state(self.active_node_idx)
220 }
221
222 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 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 pub fn initialize_node_states(&mut self, node_count: usize) {
243 self.node_states = (0..node_count).map(|_| NodeState::default()).collect();
244 }
245
246 pub fn current_block_info(&self) -> Option<BlockInfo> {
248 self.active_node_state().ok()?.current_block_info
249 }
250
251 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#[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 pub fn new() -> Self {
284 Self::default()
285 }
286
287 pub fn with_setup(mut self, setup: Setup<I>) -> Self {
289 self.setup = Some(setup);
290 self
291 }
292
293 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 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 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 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 drop(setup);
346
347 Ok(())
348 }
349}