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

reth_e2e_test_utils/
setup_builder.rs
