reth_e2e_test_utils/
setup_builder.rs

1//! Builder for configuring and launching 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::{
7    eth_payload_attributes, node::NodeTestContext, test_chain_spec, wallet::Wallet, Adapter,
8    NodeBuilderHelper, NodeHelperType, TmpNodeAdapter,
9};
10use eyre::ensure;
11use futures_util::future::{BoxFuture, TryJoinAll};
12use reth_chainspec::{ChainSpec, EthChainSpec, EthereumHardfork};
13use reth_node_api::{PayloadAttrTy, TreeConfig};
14use reth_node_builder::{
15    DebugNode, DebugNodeLauncher, EngineNodeLauncher, Node, NodeBuilder, NodeBuilderWithComponents,
16    NodeConfig, NodeHandle,
17};
18use reth_node_core::{
19    args::{DatadirArgs, DiscoveryArgs, NetworkArgs, PruningArgs, RpcServerArgs, StorageArgs},
20    dirs::{ChainPath, DataDirPath, MaybePlatformPath},
21};
22use reth_primitives_traits::AlloyBlockHeader;
23use reth_provider::providers::BlockchainProvider;
24use reth_rpc_server_types::RpcModuleSelection;
25use reth_tasks::Runtime;
26use std::{path::PathBuf, sync::Arc, time::Duration};
27use tracing::{span, Instrument, Level};
28
29/// Builder for configuring and launching test node setups.
30///
31/// By default, the nodes:
32/// - run on a shared [`Runtime::test`] runtime,
33/// - build payloads with [`eth_payload_attributes`] for the hardforks active in the chain spec,
34/// - have discovery disabled, use unused ports and serve all RPC modules except `testing` over
35///   HTTP,
36/// - report an idle sync state from startup, so they gossip transactions before their first block,
37/// - are connected to each other.
38///
39/// Once launched, each node receives a forkchoice update that makes genesis the head, safe and
40/// finalized block, unless [dev mining](Self::with_dev_mining) is enabled.
41///
42/// Configuration and tree configuration modifiers are applied in the order they are added.
43///
44/// Use [`E2ETestSetupExt::test_setup_for`] to set up nodes on the [`test_chain_spec`] at a
45/// hardfork, or [`E2ETestSetupExt::test_setup`] for any other chain spec:
46///
47/// ```ignore
48/// let (mut node, wallet) = EthereumNode::test_setup_for(EthereumHardfork::Cancun)
49///     .with_tree_config_modifier(|config| config.with_persistence_threshold(0))
50///     .build_single()
51///     .await?;
52/// ```
53pub struct E2ETestSetupBuilder<N: NodeBuilderHelper> {
54    num_nodes: usize,
55    chain_spec: Arc<N::ChainSpec>,
56    runtime: Option<Runtime>,
57    attributes_generator: Option<AttributesGenerator<N>>,
58    connect_nodes: bool,
59    tree_config_modifiers: Vec<TreeConfigModifier>,
60    node_config_modifiers: Vec<NodeConfigModifier<N::ChainSpec>>,
61    storage_v2: bool,
62    dev_launcher: Option<NodeLauncher<N>>,
63    dev_payload_attributes: Option<PayloadAttributesMapper<N>>,
64}
65
66impl<N: NodeBuilderHelper> E2ETestSetupBuilder<N> {
67    /// Creates a new builder for `num_nodes` nodes of the given chain.
68    pub fn new(num_nodes: usize, chain_spec: Arc<N::ChainSpec>) -> Self {
69        Self {
70            num_nodes,
71            chain_spec,
72            runtime: None,
73            attributes_generator: None,
74            connect_nodes: true,
75            tree_config_modifiers: Vec::new(),
76            node_config_modifiers: Vec::new(),
77            storage_v2: StorageArgs::default().v2,
78            dev_launcher: None,
79            dev_payload_attributes: None,
80        }
81    }
82
83    /// Sets the number of nodes to launch.
84    pub const fn with_num_nodes(mut self, num_nodes: usize) -> Self {
85        self.num_nodes = num_nodes;
86        self
87    }
88
89    /// Launches the nodes on the given runtime instead of a new [`Runtime::test`].
90    ///
91    /// This lets multiple setups, or other components of a test, share the same tokio handle and
92    /// rayon pools. Note that the tasks of the nodes are only shut down once all handles to the
93    /// runtime are dropped, not when the nodes are dropped.
94    pub fn with_runtime(mut self, runtime: Runtime) -> Self {
95        self.runtime = Some(runtime);
96        self
97    }
98
99    /// Sets the generator for the payload attributes of the payloads built by the test nodes.
100    ///
101    /// The generator is called with the timestamp of the next payload. Defaults to
102    /// [`eth_payload_attributes`] for the chain spec of the setup.
103    pub fn with_attributes_generator<G>(mut self, generator: G) -> Self
104    where
105        G: Fn(u64) -> PayloadAttrTy<N> + Send + Sync + 'static,
106    {
107        self.attributes_generator = Some(Arc::new(generator));
108        self
109    }
110
111    /// Sets whether nodes should be interconnected (default: true).
112    pub const fn with_connect_nodes(mut self, connect_nodes: bool) -> Self {
113        self.connect_nodes = connect_nodes;
114        self
115    }
116
117    /// Adds a modifier for the tree configuration.
118    ///
119    /// The closure receives the current tree config and returns a modified version. The base
120    /// config is the default config with a small cross block cache.
121    pub fn with_tree_config_modifier<G>(mut self, modifier: G) -> Self
122    where
123        G: Fn(TreeConfig) -> TreeConfig + Send + Sync + 'static,
124    {
125        self.tree_config_modifiers.push(Box::new(modifier));
126        self
127    }
128
129    /// Adds a modifier for the node configuration.
130    ///
131    /// The closure receives the current node config and returns a modified version.
132    pub fn with_node_config_modifier<G>(mut self, modifier: G) -> Self
133    where
134        G: Fn(NodeConfig<N::ChainSpec>) -> NodeConfig<N::ChainSpec> + Send + Sync + 'static,
135    {
136        self.node_config_modifiers.push(Box::new(modifier));
137        self
138    }
139
140    /// Adds a modifier for the RPC server arguments.
141    ///
142    /// The closure receives the current arguments, which serve all modules except `testing` over
143    /// HTTP on an unused port.
144    pub fn with_rpc_modifier<G>(self, modifier: G) -> Self
145    where
146        G: Fn(RpcServerArgs) -> RpcServerArgs + Send + Sync + 'static,
147    {
148        self.with_node_config_modifier(move |mut config| {
149            config.rpc = modifier(config.rpc);
150            config
151        })
152    }
153
154    /// Sets the pruning arguments for the test nodes.
155    pub fn with_pruning(self, pruning: PruningArgs) -> Self {
156        self.with_node_config_modifier(move |config| config.with_pruning(pruning.clone()))
157    }
158
159    /// Sets whether nodes use the v2 storage layout (`--storage.v2`), which routes tx hashes,
160    /// history indices, etc. to `RocksDB` and changesets/senders to static files.
161    ///
162    /// Defaults to the node's `--storage.v2` default. Node config modifiers run afterwards and can
163    /// still override it.
164    pub const fn with_storage_v2(mut self, storage_v2: bool) -> Self {
165        self.storage_v2 = storage_v2;
166        self
167    }
168
169    /// Sets whether the nodes run in dev mode (`--dev`).
170    ///
171    /// Unlike [dev mining](Self::with_dev_mining), this only sets the dev flag and does not start a
172    /// local miner.
173    pub fn with_dev_mode(self, dev: bool) -> Self {
174        self.with_node_config_modifier(move |config| config.set_dev(dev))
175    }
176
177    /// Launches the node in dev mode with a local miner that builds a block every `block_time`, or
178    /// as soon as a transaction is pending if `None`.
179    ///
180    /// Dev mining is limited to a single node, since every node would mine its own chain. The local
181    /// miner drives forkchoice and payload building, so the node does not receive the initial
182    /// forkchoice update to genesis, and the block producing helpers of [`NodeTestContext`], e.g.
183    /// [`NodeTestContext::advance_block`] and [`NodeTestContext::update_forkchoice`], must not be
184    /// used. Use [`Self::map_dev_payload_attributes`] to customize the payload attributes of the
185    /// mined blocks.
186    pub fn with_dev_mining(mut self, block_time: Option<Duration>) -> Self
187    where
188        N: DebugNode<Adapter<N>>,
189    {
190        self.dev_launcher = Some(|args| Box::pin(launch_dev_node::<N>(args)));
191        self.with_node_config_modifier(move |mut config| {
192            config.dev.dev = true;
193            config.dev.block_time = block_time;
194            config
195        })
196    }
197
198    /// Maps the payload attributes of the blocks built by the local miner of
199    /// [dev mining](Self::with_dev_mining) nodes, e.g. to set the fee recipient.
200    ///
201    /// Mappers are applied in the order they are added. Has no effect unless dev mining is
202    /// enabled.
203    pub fn map_dev_payload_attributes<G>(mut self, map: G) -> Self
204    where
205        G: Fn(PayloadAttrTy<N>) -> PayloadAttrTy<N> + Send + Sync + 'static,
206    {
207        self.dev_payload_attributes = Some(match self.dev_payload_attributes.take() {
208            Some(prev) => Arc::new(move |attributes| map(prev(attributes))),
209            None => Arc::new(map),
210        });
211        self
212    }
213
214    /// Builds and launches the test nodes.
215    pub async fn build(self) -> eyre::Result<(Vec<NodeHelperType<N>>, Wallet)> {
216        ensure!(
217            self.dev_launcher.is_none() || self.num_nodes == 1,
218            "dev mining requires a single node setup, got {} nodes",
219            self.num_nodes
220        );
221        let dev_mining = self.dev_launcher.is_some();
222        let launch = self.dev_launcher.unwrap_or(|args| Box::pin(launch_test_node::<N>(args)));
223        let runtime = self.runtime.unwrap_or_else(Runtime::test);
224        let attributes_generator = self.attributes_generator.unwrap_or_else(|| {
225            let chain_spec = self.chain_spec.clone();
226            Arc::new(move |timestamp| eth_payload_attributes(&chain_spec, timestamp).into())
227        });
228        let tree_config = self
229            .tree_config_modifiers
230            .iter()
231            .fold(test_tree_config(), |config, modifier| modifier(config));
232
233        let mut nodes = (0..self.num_nodes)
234            .map(async |idx| {
235                let node_config = self.node_config_modifiers.iter().fold(
236                    test_node_config(self.chain_spec.clone())
237                        .with_storage(StorageArgs { v2: self.storage_v2 }),
238                    |config, modifier| modifier(config),
239                );
240                // The local miner of dev nodes drives forkchoice, unless a modifier disabled dev
241                // mode.
242                let mines = dev_mining && node_config.dev.dev;
243                let node = launch(LaunchArgs {
244                    node_config,
245                    runtime: runtime.clone(),
246                    tree_config: tree_config.clone(),
247                    datadir: reth_db::test_utils::tempdir_path(),
248                    attributes_generator: attributes_generator.clone(),
249                    dev_payload_attributes: self.dev_payload_attributes.clone(),
250                })
251                .instrument(span!(Level::INFO, "node", idx))
252                .await?;
253
254                if !mines {
255                    let genesis = node.block_hash(self.chain_spec.genesis_header().number());
256                    node.update_forkchoice(genesis, genesis).await?;
257                }
258
259                eyre::Ok(node)
260            })
261            .collect::<TryJoinAll<_>>()
262            .await?;
263
264        if self.connect_nodes {
265            for idx in 1..self.num_nodes {
266                let (prev, current) = nodes.split_at_mut(idx);
267                prev[idx - 1].connect(&mut current[0]).await;
268            }
269
270            // Connect the last node with the first if there are more than two.
271            if self.num_nodes > 2 {
272                let (first, rest) = nodes.split_at_mut(1);
273                rest.last_mut().unwrap().connect(&mut first[0]).await;
274            }
275        }
276
277        Ok((nodes, Wallet::default().with_chain_id(self.chain_spec.chain().into())))
278    }
279
280    /// Builds and launches a single test node.
281    ///
282    /// Returns an error if the builder was not configured with exactly one node.
283    pub async fn build_single(self) -> eyre::Result<(NodeHelperType<N>, Wallet)> {
284        ensure!(self.num_nodes == 1, "expected a single node setup, got {} nodes", self.num_nodes);
285        let (mut nodes, wallet) = self.build().await?;
286        Ok((nodes.pop().expect("one node was launched"), wallet))
287    }
288}
289
290impl<N: NodeBuilderHelper> std::fmt::Debug for E2ETestSetupBuilder<N> {
291    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
292        f.debug_struct("E2ETestSetupBuilder")
293            .field("num_nodes", &self.num_nodes)
294            .field("runtime", &self.runtime)
295            .field("connect_nodes", &self.connect_nodes)
296            .field("tree_config_modifiers", &self.tree_config_modifiers.len())
297            .field("node_config_modifiers", &self.node_config_modifiers.len())
298            .field("storage_v2", &self.storage_v2)
299            .field("dev_mining", &self.dev_launcher.is_some())
300            .finish_non_exhaustive()
301    }
302}
303
304/// Extension trait to create an [`E2ETestSetupBuilder`] from a node type.
305pub trait E2ETestSetupExt: NodeBuilderHelper {
306    /// Returns an [`E2ETestSetupBuilder`] for `num_nodes` nodes of this type.
307    fn test_setup(num_nodes: usize, chain_spec: Arc<Self::ChainSpec>) -> E2ETestSetupBuilder<Self> {
308        E2ETestSetupBuilder::new(num_nodes, chain_spec)
309    }
310
311    /// Returns an [`E2ETestSetupBuilder`] for a single node of this type on the
312    /// [`test_chain_spec`] with every hardfork up to and including `fork` active at genesis.
313    ///
314    /// Use [`E2ETestSetupBuilder::with_num_nodes`] to launch more nodes.
315    fn test_setup_for(fork: EthereumHardfork) -> E2ETestSetupBuilder<Self>
316    where
317        Self::ChainSpec: From<ChainSpec>,
318    {
319        let chain_spec = Arc::unwrap_or_clone(test_chain_spec(fork));
320        E2ETestSetupBuilder::new(1, Arc::new(chain_spec.into()))
321    }
322}
323
324impl<N: NodeBuilderHelper> E2ETestSetupExt for N {}
325
326/// Closure that modifies the tree configuration of the test nodes.
327type TreeConfigModifier = Box<dyn Fn(TreeConfig) -> TreeConfig + Send + Sync>;
328
329/// Closure that modifies the node configuration of each test node.
330type NodeConfigModifier<C> = Box<dyn Fn(NodeConfig<C>) -> NodeConfig<C> + Send + Sync>;
331
332/// Closure that generates the payload attributes for a given timestamp.
333pub(crate) type AttributesGenerator<N> = Arc<dyn Fn(u64) -> PayloadAttrTy<N> + Send + Sync>;
334
335/// Closure that maps payload attributes.
336type PayloadAttributesMapper<N> = Arc<dyn Fn(PayloadAttrTy<N>) -> PayloadAttrTy<N> + Send + Sync>;
337
338/// Builder of a test node that is ready to be launched.
339type TestNodeBuilder<N> = NodeBuilderWithComponents<
340    TmpNodeAdapter<N>,
341    <N as Node<TmpNodeAdapter<N>>>::ComponentsBuilder,
342    <N as Node<TmpNodeAdapter<N>>>::AddOns,
343>;
344
345/// Function that launches a single test node.
346type NodeLauncher<N> = fn(LaunchArgs<N>) -> BoxFuture<'static, eyre::Result<NodeHelperType<N>>>;
347
348/// Arguments for launching a single test node.
349pub(crate) struct LaunchArgs<N: NodeBuilderHelper> {
350    /// The node configuration.
351    pub(crate) node_config: NodeConfig<N::ChainSpec>,
352    /// The runtime to launch the node on.
353    pub(crate) runtime: Runtime,
354    /// The engine tree configuration.
355    pub(crate) tree_config: TreeConfig,
356    /// The datadir of the node.
357    pub(crate) datadir: PathBuf,
358    /// Generator for the payload attributes of the payloads built by the test context.
359    pub(crate) attributes_generator: AttributesGenerator<N>,
360    /// Mapper for the payload attributes of the local miner in dev mode.
361    pub(crate) dev_payload_attributes: Option<PayloadAttributesMapper<N>>,
362}
363
364/// Returns the base tree configuration of test nodes.
365pub(crate) fn test_tree_config() -> TreeConfig {
366    TreeConfig::default().with_cross_block_cache_size(1024 * 1024)
367}
368
369/// Returns the base configuration of a test node.
370///
371/// Discovery is disabled, all ports are unused, all RPC modules except `testing` are served over
372/// HTTP and the node reports an idle sync state from startup.
373pub(crate) fn test_node_config<C>(chain_spec: Arc<C>) -> NodeConfig<C> {
374    let mut config = NodeConfig::new(chain_spec)
375        .with_network(NetworkArgs {
376            discovery: DiscoveryArgs { disable_discovery: true, ..DiscoveryArgs::default() },
377            ..NetworkArgs::default()
378        })
379        .with_unused_ports()
380        .with_rpc(
381            RpcServerArgs::default()
382                .with_unused_ports()
383                .with_http()
384                .with_http_api(RpcModuleSelection::All),
385        );
386    // Nodes otherwise report that they are syncing until their first canonical block, which
387    // e.g. stops transaction gossip.
388    config.debug.startup_sync_state_idle = true;
389    config
390}
391
392/// Launches a test node with the engine launcher.
393pub(crate) async fn launch_test_node<N: NodeBuilderHelper>(
394    args: LaunchArgs<N>,
395) -> eyre::Result<NodeHelperType<N>> {
396    let LaunchArgs { node_config, runtime, tree_config, datadir, attributes_generator, .. } = args;
397    let (builder, datadir) = test_node_builder::<N>(node_config, datadir);
398    let NodeHandle { node, node_exit_future: _ } =
399        builder.launch_with(EngineNodeLauncher::new(runtime, datadir, tree_config)).await?;
400
401    NodeTestContext::new(node, move |timestamp| attributes_generator(timestamp)).await
402}
403
404/// Launches a test node with the debug launcher, which runs a local miner in dev mode.
405async fn launch_dev_node<N>(args: LaunchArgs<N>) -> eyre::Result<NodeHelperType<N>>
406where
407    N: NodeBuilderHelper + DebugNode<Adapter<N>>,
408{
409    let LaunchArgs {
410        node_config,
411        runtime,
412        tree_config,
413        datadir,
414        attributes_generator,
415        dev_payload_attributes,
416    } = args;
417    let (builder, datadir) = test_node_builder::<N>(node_config, datadir);
418    let launch = builder.launch_with(DebugNodeLauncher::new(EngineNodeLauncher::new(
419        runtime,
420        datadir,
421        tree_config,
422    )));
423    let launch = match dev_payload_attributes {
424        Some(map) => launch.map_debug_payload_attributes(move |attributes| map(attributes)),
425        None => launch,
426    };
427    let NodeHandle { node, node_exit_future: _ } = launch.await?;
428
429    NodeTestContext::new(node, move |timestamp| attributes_generator(timestamp)).await
430}
431
432/// Returns the builder of a test node with a temporary database in `datadir`, and the resolved
433/// datadir of the node.
434///
435/// The datadir is removed when the node is dropped.
436fn test_node_builder<N: NodeBuilderHelper>(
437    node_config: NodeConfig<N::ChainSpec>,
438    datadir: PathBuf,
439) -> (TestNodeBuilder<N>, ChainPath<DataDirPath>) {
440    let datadir_args =
441        DatadirArgs { datadir: MaybePlatformPath::from(datadir), ..node_config.datadir.clone() };
442    let node_config = node_config.with_datadir_args(datadir_args);
443    let datadir = node_config.datadir();
444    let database = reth_db::test_utils::create_test_rw_db_with_datadir(datadir.data_dir());
445    let node = N::default();
446    let builder = NodeBuilder::new(node_config)
447        .with_database(database)
448        .with_types_and_provider::<N, BlockchainProvider<_>>()
449        .with_components(node.components_builder())
450        .with_add_ons(node.add_ons());
451    (builder, datadir)
452}
reth_e2e_test_utils/setup_builder.rs

reth_e2e_test_utils/
setup_builder.rs
