From f4a2778e02869f2af1e5cf51e01d62b56018b29b Mon Sep 17 00:00:00 2001
From: fryorcraken <commits@fryorcraken.xyz>
Date: Mon, 28 Jul 2025 11:19:25 +1000
Subject: [PATCH] test: return first peer without traversing branch

---
 packages/discovery/src/dns/dns.spec.ts | 34 ++++++++++++++++++++++++--
 1 file changed, 32 insertions(+), 2 deletions(-)

diff --git a/packages/discovery/src/dns/dns.spec.ts b/packages/discovery/src/dns/dns.spec.ts
index 2d52a409d3..ce4e15c4ba 100644
--- a/packages/discovery/src/dns/dns.spec.ts
+++ b/packages/discovery/src/dns/dns.spec.ts
@@ -33,10 +33,12 @@ const errorBranchB = `enrtree-branch:${branchDomainD}`;
 class MockDNS implements DnsClient {
   private fqdnRes: Map<string, string[]>;
   private fqdnThrows: string[];
+  public hasThrown: boolean;
 
   public constructor() {
     this.fqdnRes = new Map();
     this.fqdnThrows = [];
+    this.hasThrown = false;
   }
 
   public addRes(fqdn: string, res: string[]): void {
@@ -48,11 +50,18 @@ class MockDNS implements DnsClient {
   }
 
   public resolveTXT(fqdn: string): Promise<string[]> {
-    if (this.fqdnThrows.includes(fqdn)) throw "Mock DNS throws.";
+    if (this.fqdnThrows.includes(fqdn)) {
+      this.hasThrown = true;
+      console.log("throwing");
+      throw "Mock DNS throws.";
+    }
 
     const res = this.fqdnRes.get(fqdn);
 
-    if (!res) throw `Mock DNS could not resolve ${fqdn}`;
+    if (!res) {
+      this.hasThrown = true;
+      throw `Mock DNS could not resolve ${fqdn}`;
+    }
 
     return Promise.resolve(res);
   }
@@ -163,6 +172,7 @@ describe("DNS Node Discovery", () => {
     }
     expect(peersB.length).to.eq(1);
     expect(peersA[0].ip).to.eq(peersB[0].ip);
+    expect(mockDns.hasThrown).to.be.false;
   });
 });
 
@@ -204,6 +214,26 @@ describe("DNS Node Discovery w/ capabilities", () => {
     );
   });
 
+  it("return first retrieved peers without further DNS queries", async function () {
+    mockDns.addRes(`${rootDomain}.${host}`, multiComponentBranch);
+    mockDns.addRes(`${branchDomainA}.${host}`, [
+      mockData.enrWithWaku2RelayStore
+    ]);
+    // The ENR Tree is such as there are more branches to be explored.
+    // But they should not be explored if it isn't asked
+    mockDns.addThrow(`${branchDomainB}.${host}`);
+
+    const dnsNodeDiscovery = new DnsNodeDiscovery(mockDns);
+
+    const iterator = dnsNodeDiscovery.getNextPeer([mockData.enrTree]);
+    const { value: peer } = await iterator.next();
+
+    expect(peer.peerId?.toString()).to.eq(
+      "16Uiu2HAm2HyS6brcCspSbszG9i36re2bWBVjMe3tMdnFp1Hua34F"
+    );
+    expect(mockDns.hasThrown).to.be.false;
+  });
+
   it("retrieves all peers (3) when branch entries are composed of multiple strings", async function () {
     mockDns.addRes(`${rootDomain}.${host}`, multiComponentBranch);
     mockDns.addRes(`${branchDomainA}.${host}`, [