Skip to content

Commit 01820fd

Browse files
fgimenezfgimenez
authored
feat(e2e): add builder API for configuring test node setups (#19146)
1 parent 93b63bc commit 01820fd

File tree

3 files changed

+265
-123
lines changed

3 files changed

+265
-123
lines changed

‎crates/e2e-test-utils/src/lib.rs‎

Lines changed: 18 additions & 123 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,19 @@
11
//! Utilities for end-to-end tests.
22
33
use node::NodeTestContext;
4-
use reth_chainspec::{ChainSpec, EthChainSpec};
4+
use reth_chainspec::ChainSpec;
55
use reth_db::{test_utils::TempDatabase, DatabaseEnv};
66
use reth_engine_local::LocalPayloadAttributesBuilder;
77
use reth_network_api::test_utils::PeersHandleProvider;
88
use reth_node_builder::{
99
components::NodeComponentsBuilder,
1010
rpc::{EngineValidatorAddOn, RethRpcAddOns},
11-
EngineNodeLauncher, FullNodeTypesAdapter, Node, NodeAdapter, NodeBuilder, NodeComponents,
12-
NodeConfig, NodeHandle, NodePrimitives, NodeTypes, NodeTypesWithDBAdapter,
13-
PayloadAttributesBuilder, PayloadTypes,
11+
FullNodeTypesAdapter, Node, NodeAdapter, NodeComponents, NodePrimitives, NodeTypes,
12+
NodeTypesWithDBAdapter, PayloadAttributesBuilder, PayloadTypes,
1413
};
15-
use reth_node_core::args::{DiscoveryArgs, NetworkArgs, RpcServerArgs};
1614
use reth_provider::providers::{BlockchainProvider, NodeTypesForProvider};
17-
use reth_rpc_server_types::RpcModuleSelection;
1815
use reth_tasks::TaskManager;
1916
use std::sync::Arc;
20-
use tracing::{span, Level};
2117
use wallet::Wallet;
2218

2319
/// Wrapper type to create test nodes
@@ -45,6 +41,10 @@ mod rpc;
4541
/// Utilities for creating and writing RLP test data
4642
pub mod test_rlp_utils;
4743

44+
/// Builder for configuring test node setups
45+
mod setup_builder;
46+
pub use setup_builder::E2ETestSetupBuilder;
47+
4848
/// Creates the initial setup with `num_nodes` started and interconnected.
4949
pub async fn setup<N>(
5050
num_nodes: usize,
@@ -53,60 +53,14 @@ pub async fn setup<N>(
5353
attributes_generator: impl Fn(u64) -> <<N as NodeTypes>::Payload as PayloadTypes>::PayloadBuilderAttributes + Send + Sync + Copy + 'static,
5454
) -> eyre::Result<(Vec<NodeHelperType<N>>, TaskManager, Wallet)>
5555
where
56-
N: Default + Node<TmpNodeAdapter<N>> + NodeTypesForProvider,
57-
N::ComponentsBuilder: NodeComponentsBuilder<
58-
TmpNodeAdapter<N>,
59-
Components: NodeComponents<TmpNodeAdapter<N>, Network: PeersHandleProvider>,
60-
>,
61-
N::AddOns: RethRpcAddOns<Adapter<N>> + EngineValidatorAddOn<Adapter<N>>,
56+
N: NodeBuilderHelper,
6257
LocalPayloadAttributesBuilder<N::ChainSpec>:
6358
PayloadAttributesBuilder<<<N as NodeTypes>::Payload as PayloadTypes>::PayloadAttributes>,
6459
{
65-
let tasks = TaskManager::current();
66-
let exec = tasks.executor();
67-
68-
let network_config = NetworkArgs {
69-
discovery: DiscoveryArgs { disable_discovery: true, ..DiscoveryArgs::default() },
70-
..NetworkArgs::default()
71-
};
72-
73-
// Create nodes and peer them
74-
let mut nodes: Vec<NodeTestContext<_, _>> = Vec::with_capacity(num_nodes);
75-
76-
for idx in 0..num_nodes {
77-
let node_config = NodeConfig::new(chain_spec.clone())
78-
.with_network(network_config.clone())
79-
.with_unused_ports()
80-
.with_rpc(RpcServerArgs::default().with_unused_ports().with_http())
81-
.set_dev(is_dev);
82-
83-
let span = span!(Level::INFO, "node", idx);
84-
let _enter = span.enter();
85-
let NodeHandle { node, node_exit_future: _ } = NodeBuilder::new(node_config.clone())
86-
.testing_node(exec.clone())
87-
.node(Default::default())
88-
.launch()
89-
.await?;
90-
91-
let mut node = NodeTestContext::new(node, attributes_generator).await?;
92-
93-
// Connect each node in a chain.
94-
if let Some(previous_node) = nodes.last_mut() {
95-
previous_node.connect(&mut node).await;
96-
}
97-
98-
// Connect last node with the first if there are more than two
99-
if idx + 1 == num_nodes &&
100-
num_nodes > 2 &&
101-
let Some(first_node) = nodes.first_mut()
102-
{
103-
node.connect(first_node).await;
104-
}
105-
106-
nodes.push(node);
107-
}
108-
109-
Ok((nodes, tasks, Wallet::default().with_chain_id(chain_spec.chain().into())))
60+
E2ETestSetupBuilder::new(num_nodes, chain_spec, attributes_generator)
61+
.with_node_config_modifier(move |config| config.set_dev(is_dev))
62+
.build()
63+
.await
11064
}
11165

11266
/// Creates the initial setup with `num_nodes` started and interconnected.
@@ -155,71 +109,12 @@ where
155109
LocalPayloadAttributesBuilder<N::ChainSpec>:
156110
PayloadAttributesBuilder<<N::Payload as PayloadTypes>::PayloadAttributes>,
157111
{
158-
let tasks = TaskManager::current();
159-
let exec = tasks.executor();
160-
161-
let network_config = NetworkArgs {
162-
discovery: DiscoveryArgs { disable_discovery: true, ..DiscoveryArgs::default() },
163-
..NetworkArgs::default()
164-
};
165-
166-
// Create nodes and peer them
167-
let mut nodes: Vec<NodeTestContext<_, _>> = Vec::with_capacity(num_nodes);
168-
169-
for idx in 0..num_nodes {
170-
let node_config = NodeConfig::new(chain_spec.clone())
171-
.with_network(network_config.clone())
172-
.with_unused_ports()
173-
.with_rpc(
174-
RpcServerArgs::default()
175-
.with_unused_ports()
176-
.with_http()
177-
.with_http_api(RpcModuleSelection::All),
178-
)
179-
.set_dev(is_dev);
180-
181-
let span = span!(Level::INFO, "node", idx);
182-
let _enter = span.enter();
183-
let node = N::default();
184-
let NodeHandle { node, node_exit_future: _ } = NodeBuilder::new(node_config.clone())
185-
.testing_node(exec.clone())
186-
.with_types_and_provider::<N, BlockchainProvider<_>>()
187-
.with_components(node.components_builder())
188-
.with_add_ons(node.add_ons())
189-
.launch_with_fn(|builder| {
190-
let launcher = EngineNodeLauncher::new(
191-
builder.task_executor().clone(),
192-
builder.config().datadir(),
193-
tree_config.clone(),
194-
);
195-
builder.launch_with(launcher)
196-
})
197-
.await?;
198-
199-
let mut node = NodeTestContext::new(node, attributes_generator).await?;
200-
201-
let genesis = node.block_hash(0);
202-
node.update_forkchoice(genesis, genesis).await?;
203-
204-
// Connect each node in a chain if requested.
205-
if connect_nodes {
206-
if let Some(previous_node) = nodes.last_mut() {
207-
previous_node.connect(&mut node).await;
208-
}
209-
210-
// Connect last node with the first if there are more than two
211-
if idx + 1 == num_nodes &&
212-
num_nodes > 2 &&
213-
let Some(first_node) = nodes.first_mut()
214-
{
215-
node.connect(first_node).await;
216-
}
217-
}
218-
219-
nodes.push(node);
220-
}
221-
222-
Ok((nodes, tasks, Wallet::default().with_chain_id(chain_spec.chain().into())))
112+
E2ETestSetupBuilder::new(num_nodes, chain_spec, attributes_generator)
113+
.with_tree_config_modifier(move |_| tree_config.clone())
114+
.with_node_config_modifier(move |config| config.set_dev(is_dev))
115+
.with_connect_nodes(connect_nodes)
116+
.build()
117+
.await
223118
}
224119

225120
// Type aliases

‎crates/e2e-test-utils/src/setup_builder.rs‎

Lines changed: 210 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,210 @@
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+
6+
use crate::{node::NodeTestContext, wallet::Wallet, NodeBuilderHelper, NodeHelperType, TmpDB};
7+
use reth_chainspec::EthChainSpec;
8+
use reth_engine_local::LocalPayloadAttributesBuilder;
9+
use reth_node_builder::{
10+
EngineNodeLauncher, NodeBuilder, NodeConfig, NodeHandle, NodeTypes, NodeTypesWithDBAdapter,
11+
PayloadAttributesBuilder, PayloadTypes,
12+
};
13+
use reth_node_core::args::{DiscoveryArgs, NetworkArgs, RpcServerArgs};
14+
use reth_provider::providers::BlockchainProvider;
15+
use reth_rpc_server_types::RpcModuleSelection;
16+
use reth_tasks::TaskManager;
17+
use std::sync::Arc;
18+
use tracing::{span, Level};
19+
20+
/// Type alias for tree config modifier closure
21+
type TreeConfigModifier =
22+
Box<dyn Fn(reth_node_api::TreeConfig) -> reth_node_api::TreeConfig + Send + Sync>;
23+
24+
/// Type alias for node config modifier closure
25+
type 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.
32+
pub struct E2ETestSetupBuilder<N, F>
33+
where
34+
N: NodeBuilderHelper,
35+
F: Fn(u64) -> <<N as NodeTypes>::Payload as PayloadTypes>::PayloadBuilderAttributes
36+
+ Send
37+
+ Sync
38+
+ Copy
39+
+ 'static,
40+
LocalPayloadAttributesBuilder<N::ChainSpec>:
41+
PayloadAttributesBuilder<<N::Payload as PayloadTypes>::PayloadAttributes>,
42+
{
43+
num_nodes: usize,
44+
chain_spec: Arc<N::ChainSpec>,
45+
attributes_generator: F,
46+
connect_nodes: bool,
47+
tree_config_modifier: Option<TreeConfigModifier>,
48+
node_config_modifier: Option<NodeConfigModifier<N::ChainSpec>>,
49+
}
50+
51+
impl<N, F> E2ETestSetupBuilder<N, F>
52+
where
53+
N: NodeBuilderHelper,
54+
F: Fn(u64) -> <<N as NodeTypes>::Payload as PayloadTypes>::PayloadBuilderAttributes
55+
+ Send
56+
+ Sync
57+
+ Copy
58+
+ 'static,
59+
LocalPayloadAttributesBuilder<N::ChainSpec>:
60+
PayloadAttributesBuilder<<N::Payload as PayloadTypes>::PayloadAttributes>,
61+
{
62+
/// Creates a new builder with the required parameters.
63+
pub fn new(num_nodes: usize, chain_spec: Arc<N::ChainSpec>, attributes_generator: F) -> Self {
64+
Self {
65+
num_nodes,
66+
chain_spec,
67+
attributes_generator,
68+
connect_nodes: true,
69+
tree_config_modifier: None,
70+
node_config_modifier: None,
71+
}
72+
}
73+
74+
/// Sets whether nodes should be interconnected (default: true).
75+
pub const fn with_connect_nodes(mut self, connect_nodes: bool) -> Self {
76+
self.connect_nodes = connect_nodes;
77+
self
78+
}
79+
80+
/// Sets a modifier function for the tree configuration.
81+
///
82+
/// The closure receives the base tree config and returns a modified version.
83+
pub fn with_tree_config_modifier<G>(mut self, modifier: G) -> Self
84+
where
85+
G: Fn(reth_node_api::TreeConfig) -> reth_node_api::TreeConfig + Send + Sync + 'static,
86+
{
87+
self.tree_config_modifier = Some(Box::new(modifier));
88+
self
89+
}
90+
91+
/// Sets a modifier function for the node configuration.
92+
///
93+
/// The closure receives the base node config and returns a modified version.
94+
pub fn with_node_config_modifier<G>(mut self, modifier: G) -> Self
95+
where
96+
G: Fn(NodeConfig<N::ChainSpec>) -> NodeConfig<N::ChainSpec> + Send + Sync + 'static,
97+
{
98+
self.node_config_modifier = Some(Box::new(modifier));
99+
self
100+
}
101+
102+
/// Builds and launches the test nodes.
103+
pub async fn build(
104+
self,
105+
) -> eyre::Result<(
106+
Vec<NodeHelperType<N, BlockchainProvider<NodeTypesWithDBAdapter<N, TmpDB>>>>,
107+
TaskManager,
108+
Wallet,
109+
)> {
110+
let tasks = TaskManager::current();
111+
let exec = tasks.executor();
112+
113+
let network_config = NetworkArgs {
114+
discovery: DiscoveryArgs { disable_discovery: true, ..DiscoveryArgs::default() },
115+
..NetworkArgs::default()
116+
};
117+
118+
// Apply tree config modifier if present
119+
let tree_config = if let Some(modifier) = self.tree_config_modifier {
120+
modifier(reth_node_api::TreeConfig::default())
121+
} else {
122+
reth_node_api::TreeConfig::default()
123+
};
124+
125+
let mut nodes: Vec<NodeTestContext<_, _>> = Vec::with_capacity(self.num_nodes);
126+
127+
for idx in 0..self.num_nodes {
128+
// Create base node config
129+
let base_config = NodeConfig::new(self.chain_spec.clone())
130+
.with_network(network_config.clone())
131+
.with_unused_ports()
132+
.with_rpc(
133+
RpcServerArgs::default()
134+
.with_unused_ports()
135+
.with_http()
136+
.with_http_api(RpcModuleSelection::All),
137+
);
138+
139+
// Apply node config modifier if present
140+
let node_config = if let Some(modifier) = &self.node_config_modifier {
141+
modifier(base_config)
142+
} else {
143+
base_config
144+
};
145+
146+
let span = span!(Level::INFO, "node", idx);
147+
let _enter = span.enter();
148+
let node = N::default();
149+
let NodeHandle { node, node_exit_future: _ } = NodeBuilder::new(node_config)
150+
.testing_node(exec.clone())
151+
.with_types_and_provider::<N, BlockchainProvider<_>>()
152+
.with_components(node.components_builder())
153+
.with_add_ons(node.add_ons())
154+
.launch_with_fn(|builder| {
155+
let launcher = EngineNodeLauncher::new(
156+
builder.task_executor().clone(),
157+
builder.config().datadir(),
158+
tree_config.clone(),
159+
);
160+
builder.launch_with(launcher)
161+
})
162+
.await?;
163+
164+
let mut node = NodeTestContext::new(node, self.attributes_generator).await?;
165+
166+
let genesis = node.block_hash(0);
167+
node.update_forkchoice(genesis, genesis).await?;
168+
169+
// Connect nodes if requested
170+
if self.connect_nodes {
171+
if let Some(previous_node) = nodes.last_mut() {
172+
previous_node.connect(&mut node).await;
173+
}
174+
175+
// Connect last node with the first if there are more than two
176+
if idx + 1 == self.num_nodes &&
177+
self.num_nodes > 2 &&
178+
let Some(first_node) = nodes.first_mut()
179+
{
180+
node.connect(first_node).await;
181+
}
182+
}
183+
184+
nodes.push(node);
185+
}
186+
187+
Ok((nodes, tasks, Wallet::default().with_chain_id(self.chain_spec.chain().into())))
188+
}
189+
}
190+
191+
impl<N, F> std::fmt::Debug for E2ETestSetupBuilder<N, F>
192+
where
193+
N: NodeBuilderHelper,
194+
F: Fn(u64) -> <<N as NodeTypes>::Payload as PayloadTypes>::PayloadBuilderAttributes
195+
+ Send
196+
+ Sync
197+
+ Copy
198+
+ 'static,
199+
LocalPayloadAttributesBuilder<N::ChainSpec>:
200+
PayloadAttributesBuilder<<N::Payload as PayloadTypes>::PayloadAttributes>,
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+
}

0 commit comments

Comments
 (0)