diff --git a/crates/net/network/src/config.rs b/crates/net/network/src/config.rs index 860dc9c929..043b4c1009 100644 --- a/crates/net/network/src/config.rs +++ b/crates/net/network/src/config.rs @@ -271,29 +271,12 @@ impl NetworkConfigBuilder { self } - /// Disables Discv4 discovery. - pub fn no_discv4_discovery(mut self) -> Self { - self.discovery_v4_builder = None; - self - } - /// Sets the dns discovery config to use. pub fn dns_discovery(mut self, config: DnsDiscoveryConfig) -> Self { self.dns_discovery_config = Some(config); self } - /// Disables DNS discovery - pub fn no_dns_discovery(mut self) -> Self { - self.dns_discovery_config = None; - self - } - - /// Disables all discovery. - pub fn no_discovery(self) -> Self { - self.no_discv4_discovery().no_dns_discovery() - } - /// Sets the boot nodes. pub fn boot_nodes(mut self, nodes: impl IntoIterator) -> Self { self.boot_nodes = nodes.into_iter().collect(); @@ -306,12 +289,44 @@ impl NetworkConfigBuilder { self } + /// Disables all discovery. + pub fn disable_discovery(self) -> Self { + self.disable_discv4_discovery().disable_dns_discovery() + } + + /// Disables all discovery if the given condition is true. + pub fn disable_discovery_if(self, disable: bool) -> Self { + if disable { + self.disable_discovery() + } else { + self + } + } + /// Disable the Discv4 discovery. pub fn disable_discv4_discovery(mut self) -> Self { self.discovery_v4_builder = None; self } + /// Disable the DNS discovery if the given condition is true. + pub fn disable_dns_discovery_if(self, disable: bool) -> Self { + if disable { + self.disable_dns_discovery() + } else { + self + } + } + + /// Disable the Discv4 discovery if the given condition is true. + pub fn disable_discv4_discovery_if(self, disable: bool) -> Self { + if disable { + self.disable_discv4_discovery() + } else { + self + } + } + /// Consumes the type and creates the actual [`NetworkConfig`] /// for the given client type that can interact with the chain. pub fn build(self, client: C) -> NetworkConfig { diff --git a/crates/net/network/src/test_utils/testnet.rs b/crates/net/network/src/test_utils/testnet.rs index bfd2645a47..190d94164b 100644 --- a/crates/net/network/src/test_utils/testnet.rs +++ b/crates/net/network/src/test_utils/testnet.rs @@ -316,8 +316,8 @@ where NetworkConfigBuilder::new(secret_key) .listener_addr(SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 0))) .discovery_addr(SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 0))) - .no_dns_discovery() - .no_discv4_discovery() + .disable_dns_discovery() + .disable_discv4_discovery() } } diff --git a/crates/net/network/src/transactions.rs b/crates/net/network/src/transactions.rs index 2847436e54..c1fdd26314 100644 --- a/crates/net/network/src/transactions.rs +++ b/crates/net/network/src/transactions.rs @@ -728,8 +728,10 @@ mod tests { let client = NoopProvider::default(); let pool = testing_pool(); - let config = - NetworkConfigBuilder::new(secret_key).no_discovery().listener_port(0).build(client); + let config = NetworkConfigBuilder::new(secret_key) + .disable_discovery() + .listener_port(0) + .build(client); let (handle, network, mut transactions, _) = NetworkManager::new(config) .await .unwrap() @@ -771,8 +773,10 @@ mod tests { let client = NoopProvider::default(); let pool = testing_pool(); - let config = - NetworkConfigBuilder::new(secret_key).no_discovery().listener_port(0).build(client); + let config = NetworkConfigBuilder::new(secret_key) + .disable_discovery() + .listener_port(0) + .build(client); let (network_handle, network, mut transactions, _) = NetworkManager::new(config) .await .unwrap() @@ -844,8 +848,10 @@ mod tests { let client = NoopProvider::default(); let pool = testing_pool(); - let config = - NetworkConfigBuilder::new(secret_key).no_discovery().listener_port(0).build(client); + let config = NetworkConfigBuilder::new(secret_key) + .disable_discovery() + .listener_port(0) + .build(client); let (network_handle, network, mut transactions, _) = NetworkManager::new(config) .await .unwrap()