reth_e2e_test_utils/testsuite/
setup.rs

Help
1//! Test setup utilities for configuring the initial state.
2
3use crate::{setup_engine, testsuite::Environment, NodeBuilderHelper, PayloadAttributesBuilder};
4use alloy_eips::BlockNumberOrTag;
5use alloy_primitives::B256;
6use alloy_rpc_types_engine::PayloadAttributes;
7use alloy_rpc_types_eth::{Block as RpcBlock, Header, Receipt, Transaction};
8use eyre::{eyre, Result};
9use reth_chainspec::ChainSpec;
10use reth_engine_local::LocalPayloadAttributesBuilder;
11use reth_ethereum_primitives::Block;
12use reth_node_api::{NodeTypes, PayloadTypes};
13use reth_node_core::primitives::RecoveredBlock;
14use reth_payload_builder::EthPayloadBuilderAttributes;
15use reth_rpc_api::clients::EthApiClient;
16use revm::state::EvmState;
17use std::{marker::PhantomData, sync::Arc};
18use tokio::{
19    sync::mpsc,
20    time::{sleep, Duration},
21};
22use tracing::{debug, error};
23
24/// Configuration for setting upa test environment
25#[derive(Debug)]
26pub struct Setup<I> {
27    /// Chain specification to use
28    pub chain_spec: Option<Arc<ChainSpec>>,
29    /// Genesis block to use
30    pub genesis: Option<Genesis>,
31    /// Blocks to replay during setup
32    pub blocks: Vec<RecoveredBlock<Block>>,
33    /// Initial state to load
34    pub state: Option<EvmState>,
35    /// Network configuration
36    pub network: NetworkSetup,
37    /// Shutdown channel to stop nodes when setup is dropped
38    shutdown_tx: Option<mpsc::Sender<()>>,
39    /// Is this setup in dev mode
40    pub is_dev: bool,
41    /// Tracks instance generic.
42    _phantom: PhantomData<I>,
43}
44
45impl<I> Default for Setup<I> {
46    fn default() -> Self {
47        Self {
48            chain_spec: None,
49            genesis: None,
50            blocks: Vec::new(),
51            state: None,
52            network: NetworkSetup::default(),
53            shutdown_tx: None,
54            is_dev: true,
55            _phantom: Default::default(),
56        }
57    }
58}
59
60impl<I> Drop for Setup<I> {
61    fn drop(&mut self) {
62        // Send shutdown signal if the channel exists
63        if let Some(tx) = self.shutdown_tx.take() {
64            let _ = tx.try_send(());
65        }
66    }
67}
68
69impl<I> Setup<I> {
70    /// Create a new setup with default values
71    pub fn new() -> Self {
72        Self::default()
73    }
74
75    /// Set the chain specification
76    pub fn with_chain_spec(mut self, chain_spec: Arc<ChainSpec>) -> Self {
77        self.chain_spec = Some(chain_spec);
78        self
79    }
80
81    /// Set the genesis block
82    pub const fn with_genesis(mut self, genesis: Genesis) -> Self {
83        self.genesis = Some(genesis);
84        self
85    }
86
87    /// Add a block to replay during setup
88    pub fn with_block(mut self, block: RecoveredBlock<Block>) -> Self {
89        self.blocks.push(block);
90        self
91    }
92
93    /// Add multiple blocks to replay during setup
94    pub fn with_blocks(mut self, blocks: Vec<RecoveredBlock<Block>>) -> Self {
95        self.blocks.extend(blocks);
96        self
97    }
98
99    /// Set the initial state
100    pub fn with_state(mut self, state: EvmState) -> Self {
101        self.state = Some(state);
102        self
103    }
104
105    /// Set the network configuration
106    pub const fn with_network(mut self, network: NetworkSetup) -> Self {
107        self.network = network;
108        self
109    }
110
111    /// Set dev mode
112    pub const fn with_dev_mode(mut self, is_dev: bool) -> Self {
113        self.is_dev = is_dev;
114        self
115    }
116
117    /// Apply the setup to the environment
118    pub async fn apply<N>(&mut self, env: &mut Environment<I>) -> Result<()>
119    where
120        N: NodeBuilderHelper,
121        LocalPayloadAttributesBuilder<N::ChainSpec>: PayloadAttributesBuilder<
122            <<N as NodeTypes>::Payload as PayloadTypes>::PayloadAttributes,
123        >,
124    {
125        let chain_spec =
126            self.chain_spec.clone().ok_or_else(|| eyre!("Chain specification is required"))?;
127
128        let (shutdown_tx, mut shutdown_rx) = mpsc::channel(1);
129
130        self.shutdown_tx = Some(shutdown_tx);
131
132        let is_dev = self.is_dev;
133        let node_count = self.network.node_count;
134
135        let attributes_generator = move |timestamp| {
136            let attributes = PayloadAttributes {
137                timestamp,
138                prev_randao: B256::ZERO,
139                suggested_fee_recipient: alloy_primitives::Address::ZERO,
140                withdrawals: Some(vec![]),
141                parent_beacon_block_root: Some(B256::ZERO),
142            };
143            <<N as NodeTypes>::Payload as PayloadTypes>::PayloadBuilderAttributes::from(
144                EthPayloadBuilderAttributes::new(B256::ZERO, attributes),
145            )
146        };
147
148        let result = setup_engine::<N>(
149            node_count,
150            Arc::<N::ChainSpec>::new((*chain_spec).clone().into()),
151            is_dev,
152            attributes_generator,
153        )
154        .await;
155
156        let mut node_clients = Vec::new();
157        match result {
158            Ok((nodes, executor, _wallet)) => {
159                // create HTTP clients for each node's RPC and Engine API endpoints
160                for node in &nodes {
161                    let rpc = node
162                        .rpc_client()
163                        .ok_or_else(|| eyre!("Failed to create HTTP RPC client for node"))?;
164                    let engine = node.engine_api_client();
165
166                    node_clients.push(crate::testsuite::NodeClient { rpc, engine });
167                }
168
169                // spawn a separate task just to handle the shutdown
170                tokio::spawn(async move {
171                    // keep nodes and executor in scope to ensure they're not dropped
172                    let _nodes = nodes;
173                    let _executor = executor;
174                    // Wait for shutdown signal
175                    let _ = shutdown_rx.recv().await;
176                    // nodes and executor will be dropped here when the test completes
177                });
178            }
179            Err(e) => {
180                error!("Failed to setup nodes: {}", e);
181                return Err(eyre!("Failed to setup nodes: {}", e));
182            }
183        }
184
185        if node_clients.is_empty() {
186            return Err(eyre!("No nodes were created"));
187        }
188
189        // wait for all nodes to be ready to accept RPC requests before proceeding
190        for (idx, client) in node_clients.iter().enumerate() {
191            let mut retry_count = 0;
192            const MAX_RETRIES: usize = 5;
193            let mut last_error = None;
194
195            while retry_count < MAX_RETRIES {
196                match EthApiClient::<Transaction, RpcBlock, Receipt, Header>::block_by_number(
197                    &client.rpc,
198                    BlockNumberOrTag::Latest,
199                    false,
200                )
201                .await
202                {
203                    Ok(_) => {
204                        debug!("Node {idx} RPC endpoint is ready");
205                        break;
206                    }
207                    Err(e) => {
208                        last_error = Some(e);
209                        retry_count += 1;
210                        debug!(
211                            "Node {idx} RPC endpoint not ready, retry {retry_count}/{MAX_RETRIES}"
212                        );
213                        sleep(Duration::from_millis(500)).await;
214                    }
215                }
216            }
217            if retry_count == MAX_RETRIES {
218                return Err(eyre!("Failed to connect to node {idx} RPC endpoint after {MAX_RETRIES} retries: {:?}", last_error));
219            }
220        }
221
222        env.node_clients = node_clients;
223
224        // TODO: For each block in self.blocks, replay it on the node
225
226        Ok(())
227    }
228}
229
230/// Genesis block configuration
231#[derive(Debug)]
232pub struct Genesis {}
233
234/// Network configuration for setup
235#[derive(Debug, Default)]
236pub struct NetworkSetup {
237    /// Number of nodes to create
238    pub node_count: usize,
239}
240
241impl NetworkSetup {
242    /// Create a new network setup with a single node
243    pub const fn single_node() -> Self {
244        Self { node_count: 1 }
245    }
246
247    /// Create a new network setup with multiple nodes
248    pub const fn multi_node(count: usize) -> Self {
249        Self { node_count: count }
250    }
251}
reth_e2e_test_utils/testsuite/setup.rs

reth_e2e_test_utils/testsuite/
setup.rs
