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

reth_e2e_test_utils/
setup_builder.rs
