reth_e2e_test_utils/
setup_builder.rs

1//! Builder for configuring and creating test node setups.
2//!
3//! This module provides a flexible builder API for setting up test nodes with custom
4//! configurations through closures that modify `NodeConfig` and `TreeConfig`.
5
6use crate::{node::NodeTestContext, wallet::Wallet, NodeBuilderHelper, NodeHelperType, TmpDB};
7use futures_util::future::TryJoinAll;
8use reth_chainspec::EthChainSpec;
9use reth_node_builder::{
10    EngineNodeLauncher, NodeBuilder, NodeConfig, NodeHandle, NodeTypes, NodeTypesWithDBAdapter,
11    PayloadTypes,
12};
13use reth_node_core::args::{DiscoveryArgs, NetworkArgs, RpcServerArgs};
14use reth_primitives_traits::AlloyBlockHeader;
15use reth_provider::providers::BlockchainProvider;
16use reth_rpc_server_types::RpcModuleSelection;
17use reth_tasks::Runtime;
18use std::sync::Arc;
19use tracing::{span, Instrument, Level};
20
21/// Type alias for tree config modifier closure
22type TreeConfigModifier =
23    Box<dyn Fn(reth_node_api::TreeConfig) -> reth_node_api::TreeConfig + Send + Sync>;
24
25/// Type alias for node config modifier closure
26type NodeConfigModifier<C> = Box<dyn Fn(NodeConfig<C>) -> NodeConfig<C> + Send + Sync>;
27
28/// Builder for configuring and creating test node setups.
29///
30/// This builder allows customizing test node configurations through closures that
31/// modify `NodeConfig` and `TreeConfig`. It avoids code duplication by centralizing
32/// the node creation logic.
33pub struct E2ETestSetupBuilder<N, F>
34where
35    N: NodeBuilderHelper,
36    F: Fn(u64) -> <<N as NodeTypes>::Payload as PayloadTypes>::PayloadBuilderAttributes
37        + Send
38        + Sync
39        + Copy
40        + 'static,
41{
42    num_nodes: usize,
43    chain_spec: Arc<N::ChainSpec>,
44    attributes_generator: F,
45    connect_nodes: bool,
46    tree_config_modifier: Option<TreeConfigModifier>,
47    node_config_modifier: Option<NodeConfigModifier<N::ChainSpec>>,
48}
49
50impl<N, F> E2ETestSetupBuilder<N, F>
51where
52    N: NodeBuilderHelper,
53    F: Fn(u64) -> <<N as NodeTypes>::Payload as PayloadTypes>::PayloadBuilderAttributes
54        + Send
55        + Sync
56        + Copy
57        + 'static,
58{
59    /// Creates a new builder with the required parameters.
60    pub fn new(num_nodes: usize, chain_spec: Arc<N::ChainSpec>, attributes_generator: F) -> Self {
61        Self {
62            num_nodes,
63            chain_spec,
64            attributes_generator,
65            connect_nodes: true,
66            tree_config_modifier: None,
67            node_config_modifier: None,
68        }
69    }
70
71    /// Sets whether nodes should be interconnected (default: true).
72    pub const fn with_connect_nodes(mut self, connect_nodes: bool) -> Self {
73        self.connect_nodes = connect_nodes;
74        self
75    }
76
77    /// Sets a modifier function for the tree configuration.
78    ///
79    /// The closure receives the base tree config and returns a modified version.
80    pub fn with_tree_config_modifier<G>(mut self, modifier: G) -> Self
81    where
82        G: Fn(reth_node_api::TreeConfig) -> reth_node_api::TreeConfig + Send + Sync + 'static,
83    {
84        self.tree_config_modifier = Some(Box::new(modifier));
85        self
86    }
87
88    /// Sets a modifier function for the node configuration.
89    ///
90    /// The closure receives the base node config and returns a modified version.
91    pub fn with_node_config_modifier<G>(mut self, modifier: G) -> Self
92    where
93        G: Fn(NodeConfig<N::ChainSpec>) -> NodeConfig<N::ChainSpec> + Send + Sync + 'static,
94    {
95        self.node_config_modifier = Some(Box::new(modifier));
96        self
97    }
98
99    /// Enables v2 storage defaults (`--storage.v2`), routing tx hashes, history
100    /// indices, etc. to `RocksDB` and changesets/senders to static files.
101    pub fn with_storage_v2(self) -> Self {
102        self.with_node_config_modifier(|mut config| {
103            config.storage.v2 = true;
104            config
105        })
106    }
107
108    /// Builds and launches the test nodes.
109    pub async fn build(
110        self,
111    ) -> eyre::Result<(
112        Vec<NodeHelperType<N, BlockchainProvider<NodeTypesWithDBAdapter<N, TmpDB>>>>,
113        Wallet,
114    )> {
115        let runtime = Runtime::with_existing_handle(tokio::runtime::Handle::current())?;
116
117        let network_config = NetworkArgs {
118            discovery: DiscoveryArgs { disable_discovery: true, ..DiscoveryArgs::default() },
119            ..NetworkArgs::default()
120        };
121
122        // Apply tree config modifier if present, with test-appropriate defaults
123        let base_tree_config =
124            reth_node_api::TreeConfig::default().with_cross_block_cache_size(1024 * 1024);
125        let tree_config = if let Some(modifier) = self.tree_config_modifier {
126            modifier(base_tree_config)
127        } else {
128            base_tree_config
129        };
130
131        let mut nodes = (0..self.num_nodes)
132            .map(async |idx| {
133                // Create base node config
134                let base_config = NodeConfig::new(self.chain_spec.clone())
135                    .with_network(network_config.clone())
136                    .with_unused_ports()
137                    .with_rpc(
138                        RpcServerArgs::default()
139                            .with_unused_ports()
140                            .with_http()
141                            .with_http_api(RpcModuleSelection::All),
142                    );
143
144                // Apply node config modifier if present
145                let node_config = if let Some(modifier) = &self.node_config_modifier {
146                    modifier(base_config)
147                } else {
148                    base_config
149                };
150
151                let span = span!(Level::INFO, "node", idx);
152                let node = N::default();
153                let NodeHandle { node, node_exit_future: _ } = NodeBuilder::new(node_config)
154                    .testing_node(runtime.clone())
155                    .with_types_and_provider::<N, BlockchainProvider<_>>()
156                    .with_components(node.components_builder())
157                    .with_add_ons(node.add_ons())
158                    .launch_with_fn(|builder| {
159                        let launcher = EngineNodeLauncher::new(
160                            builder.task_executor().clone(),
161                            builder.config().datadir(),
162                            tree_config.clone(),
163                        );
164                        builder.launch_with(launcher)
165                    })
166                    .instrument(span)
167                    .await?;
168
169                let node = NodeTestContext::new(node, self.attributes_generator).await?;
170                let genesis_number = self.chain_spec.genesis_header().number();
171                let genesis = node.block_hash(genesis_number);
172                node.update_forkchoice(genesis, genesis).await?;
173
174                eyre::Ok(node)
175            })
176            .collect::<TryJoinAll<_>>()
177            .await?;
178
179        for idx in 0..self.num_nodes {
180            let (prev, current) = nodes.split_at_mut(idx);
181            let current = current.first_mut().unwrap();
182            // Connect nodes if requested
183            if self.connect_nodes {
184                if let Some(prev_idx) = idx.checked_sub(1) {
185                    prev[prev_idx].connect(current).await;
186                }
187
188                // Connect last node with the first if there are more than two
189                if idx + 1 == self.num_nodes &&
190                    self.num_nodes > 2 &&
191                    let Some(first) = prev.first_mut()
192                {
193                    current.connect(first).await;
194                }
195            }
196        }
197
198        Ok((nodes, Wallet::default().with_chain_id(self.chain_spec.chain().into())))
199    }
200}
201
202impl<N, F> std::fmt::Debug for E2ETestSetupBuilder<N, F>
203where
204    N: NodeBuilderHelper,
205    F: Fn(u64) -> <<N as NodeTypes>::Payload as PayloadTypes>::PayloadBuilderAttributes
206        + Send
207        + Sync
208        + Copy
209        + 'static,
210{
211    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
212        f.debug_struct("E2ETestSetupBuilder")
213            .field("num_nodes", &self.num_nodes)
214            .field("connect_nodes", &self.connect_nodes)
215            .field("tree_config_modifier", &self.tree_config_modifier.as_ref().map(|_| "<closure>"))
216            .field("node_config_modifier", &self.node_config_modifier.as_ref().map(|_| "<closure>"))
217            .finish_non_exhaustive()
218    }
219}
reth_e2e_test_utils/setup_builder.rs

reth_e2e_test_utils/
setup_builder.rs
