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>::PayloadAttributes
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>::PayloadAttributes
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    /// Sets the pruning arguments for the test nodes.
100    pub fn with_pruning(self, pruning: reth_node_core::args::PruningArgs) -> Self {
101        self.with_node_config_modifier(move |config| config.with_pruning(pruning.clone()))
102    }
103
104    /// Enables v2 storage defaults (`--storage.v2`), routing tx hashes, history
105    /// indices, etc. to `RocksDB` and changesets/senders to static files.
106    pub fn with_storage_v2(self) -> Self {
107        self.with_node_config_modifier(|mut config| {
108            config.storage.v2 = true;
109            config
110        })
111    }
112
113    /// Builds and launches the test nodes.
114    pub async fn build(
115        self,
116    ) -> eyre::Result<(
117        Vec<NodeHelperType<N, BlockchainProvider<NodeTypesWithDBAdapter<N, TmpDB>>>>,
118        Wallet,
119    )> {
120        let runtime = Runtime::test();
121
122        let network_config = NetworkArgs {
123            discovery: DiscoveryArgs { disable_discovery: true, ..DiscoveryArgs::default() },
124            ..NetworkArgs::default()
125        };
126
127        // Apply tree config modifier if present, with test-appropriate defaults
128        let base_tree_config =
129            reth_node_api::TreeConfig::default().with_cross_block_cache_size(1024 * 1024);
130        let tree_config = if let Some(modifier) = self.tree_config_modifier {
131            modifier(base_tree_config)
132        } else {
133            base_tree_config
134        };
135
136        let mut nodes = (0..self.num_nodes)
137            .map(async |idx| {
138                // Create base node config
139                let base_config = NodeConfig::new(self.chain_spec.clone())
140                    .with_network(network_config.clone())
141                    .with_unused_ports()
142                    .with_rpc(
143                        RpcServerArgs::default()
144                            .with_unused_ports()
145                            .with_http()
146                            .with_http_api(RpcModuleSelection::All),
147                    );
148
149                // Apply node config modifier if present
150                let node_config = if let Some(modifier) = &self.node_config_modifier {
151                    modifier(base_config)
152                } else {
153                    base_config
154                };
155
156                let span = span!(Level::INFO, "node", idx);
157                let node = N::default();
158                let NodeHandle { node, node_exit_future: _ } = NodeBuilder::new(node_config)
159                    .testing_node(runtime.clone())
160                    .with_types_and_provider::<N, BlockchainProvider<_>>()
161                    .with_components(node.components_builder())
162                    .with_add_ons(node.add_ons())
163                    .launch_with_fn(|builder| {
164                        let launcher = EngineNodeLauncher::new(
165                            builder.task_executor().clone(),
166                            builder.config().datadir(),
167                            tree_config.clone(),
168                        );
169                        builder.launch_with(launcher)
170                    })
171                    .instrument(span)
172                    .await?;
173
174                let node = NodeTestContext::new(node, self.attributes_generator).await?;
175                let genesis_number = self.chain_spec.genesis_header().number();
176                let genesis = node.block_hash(genesis_number);
177                node.update_forkchoice(genesis, genesis).await?;
178
179                eyre::Ok(node)
180            })
181            .collect::<TryJoinAll<_>>()
182            .await?;
183
184        for idx in 0..self.num_nodes {
185            let (prev, current) = nodes.split_at_mut(idx);
186            let current = current.first_mut().unwrap();
187            // Connect nodes if requested
188            if self.connect_nodes {
189                if let Some(prev_idx) = idx.checked_sub(1) {
190                    prev[prev_idx].connect(current).await;
191                }
192
193                // Connect last node with the first if there are more than two
194                if idx + 1 == self.num_nodes &&
195                    self.num_nodes > 2 &&
196                    let Some(first) = prev.first_mut()
197                {
198                    current.connect(first).await;
199                }
200            }
201        }
202
203        Ok((nodes, Wallet::default().with_chain_id(self.chain_spec.chain().into())))
204    }
205}
206
207impl<N, F> std::fmt::Debug for E2ETestSetupBuilder<N, F>
208where
209    N: NodeBuilderHelper,
210    F: Fn(u64) -> <<N as NodeTypes>::Payload as PayloadTypes>::PayloadAttributes
211        + Send
212        + Sync
213        + Copy
214        + 'static,
215{
216    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
217        f.debug_struct("E2ETestSetupBuilder")
218            .field("num_nodes", &self.num_nodes)
219            .field("connect_nodes", &self.connect_nodes)
220            .field("tree_config_modifier", &self.tree_config_modifier.as_ref().map(|_| "<closure>"))
221            .field("node_config_modifier", &self.node_config_modifier.as_ref().map(|_| "<closure>"))
222            .finish_non_exhaustive()
223    }
224}
reth_e2e_test_utils/setup_builder.rs

reth_e2e_test_utils/
setup_builder.rs
