From 1d6abdcff330cbe1230d968915e90f8760516ea7 Mon Sep 17 00:00:00 2001
From: "Franz Heinzmann (Frando)" <frando@unbiskant.org>
Date: Fri, 3 May 2024 00:15:03 +0200
Subject: [PATCH] cleanup

---
 iroh-net/src/discovery.rs | 61 +++++++++------------------------------
 1 file changed, 13 insertions(+), 48 deletions(-)

diff --git a/iroh-net/src/discovery.rs b/iroh-net/src/discovery.rs
index 10ddd407f6a..9f11696c9bb 100644
--- a/iroh-net/src/discovery.rs
+++ b/iroh-net/src/discovery.rs
@@ -558,15 +558,13 @@ mod tests {
 /// publish to. The DNS and pkarr servers share their state.
 #[cfg(test)]
 mod test_dns_pkarr {
-    use std::net::SocketAddr;
     use std::time::Duration;
 
     use anyhow::Result;
     use iroh_base::key::SecretKey;
-    use url::Url;
 
     use crate::{
-        discovery::{dns::DnsDiscovery, pkarr_publish::PkarrPublisher, ConcurrentDiscovery},
+        discovery::{pkarr_publish::PkarrPublisher},
         dns::node_info::{lookup_by_id, NodeInfo},
         relay::{RelayMap, RelayMode},
         test_utils::{
@@ -642,26 +640,13 @@ mod test_dns_pkarr {
     async fn pkarr_publish_dns_discover() -> Result<()> {
         let _logging_guard = iroh_test::logging::setup();
 
-        let origin = "testdns.example".to_string();
         let timeout = Duration::from_secs(2);
 
-        let dns_pkarr_server = DnsPkarrServer::run_with_origin(origin.clone()).await?;
+        let dns_pkarr_server = DnsPkarrServer::run().await?;
         let (relay_map, _relay_url, _relay_guard) = run_relay_server().await?;
 
-        let ep1 = ep_with_discovery(
-            relay_map.clone(),
-            dns_pkarr_server.nameserver,
-            &origin,
-            &dns_pkarr_server.pkarr_url,
-        )
-        .await?;
-        let ep2 = ep_with_discovery(
-            relay_map,
-            dns_pkarr_server.nameserver,
-            &origin,
-            &dns_pkarr_server.pkarr_url,
-        )
-        .await?;
+        let ep1 = ep_with_discovery(&relay_map, &dns_pkarr_server).await?;
+        let ep2 = ep_with_discovery(&relay_map, &dns_pkarr_server).await?;
 
         // wait until our shared state received the update from pkarr publishing
         dns_pkarr_server.on_node(&ep1.node_id(), timeout).await?;
@@ -676,26 +661,13 @@ mod test_dns_pkarr {
     async fn pkarr_publish_dns_discover_empty_node_addr() -> Result<()> {
         let _logging_guard = iroh_test::logging::setup();
 
-        let origin = "testdns.example".to_string();
         let timeout = Duration::from_secs(2);
 
-        let dns_pkarr_server = DnsPkarrServer::run_with_origin(origin.clone()).await?;
+        let dns_pkarr_server = DnsPkarrServer::run().await?;
         let (relay_map, _relay_url, _relay_guard) = run_relay_server().await?;
 
-        let ep1 = ep_with_discovery(
-            relay_map.clone(),
-            dns_pkarr_server.nameserver,
-            &origin,
-            &dns_pkarr_server.pkarr_url,
-        )
-        .await?;
-        let ep2 = ep_with_discovery(
-            relay_map,
-            dns_pkarr_server.nameserver,
-            &origin,
-            &dns_pkarr_server.pkarr_url,
-        )
-        .await?;
+        let ep1 = ep_with_discovery(&relay_map, &dns_pkarr_server).await?;
+        let ep2 = ep_with_discovery(&relay_map, &dns_pkarr_server).await?;
 
         // wait until our shared state received the update from pkarr publishing
         dns_pkarr_server.on_node(&ep1.node_id(), timeout).await?;
@@ -712,24 +684,17 @@ mod test_dns_pkarr {
     }
 
     async fn ep_with_discovery(
-        relay_map: RelayMap,
-        nameserver: SocketAddr,
-        node_origin: &str,
-        pkarr_relay: &Url,
+        relay_map: &RelayMap,
+        dns_pkarr_server: &DnsPkarrServer,
     ) -> Result<MagicEndpoint> {
         let secret_key = SecretKey::generate();
-        let resolver = create_dns_resolver(nameserver)?;
-        let discovery = ConcurrentDiscovery::from_services(vec![
-            Box::new(DnsDiscovery::new(node_origin.to_string())),
-            Box::new(PkarrPublisher::new(secret_key.clone(), pkarr_relay.clone())),
-        ]);
         let ep = MagicEndpoint::builder()
-            .relay_mode(RelayMode::Custom(relay_map))
+            .relay_mode(RelayMode::Custom(relay_map.clone()))
             .insecure_skip_relay_cert_verify(true)
-            .secret_key(secret_key)
-            .dns_resolver(resolver)
+            .secret_key(secret_key.clone())
             .alpns(vec![TEST_ALPN.to_vec()])
-            .discovery(Box::new(discovery))
+            .dns_resolver(dns_pkarr_server.dns_resolver())
+            .discovery(dns_pkarr_server.discovery(secret_key))
             .bind(0)
             .await?;
         Ok(ep)