From 67d5a9bfe59196601252aff7ba497c3c822df183 Mon Sep 17 00:00:00 2001
From: Marco Munizaga <git@marcopolo.io>
Date: Tue, 1 Oct 2024 09:42:14 -0700
Subject: [PATCH] Make test fn resolveAllDNS better

---
 resolve_test.go | 27 ++++++++++++++++++---------
 1 file changed, 18 insertions(+), 9 deletions(-)

diff --git a/resolve_test.go b/resolve_test.go
index f09b939..7b048eb 100644
--- a/resolve_test.go
+++ b/resolve_test.go
@@ -123,19 +123,28 @@ func TestResolveOnlyOnce(t *testing.T) {
 }
 
 func resolveAllDNS(ctx context.Context, resolver *Resolver, in ma.Multiaddr) ([]ma.Multiaddr, error) {
-	var inAddrs []ma.Multiaddr
-	outAddrs := []ma.Multiaddr{in}
-
-	for len(inAddrs) != len(outAddrs) {
-		inAddrs = outAddrs
-		outAddrs = nil
-		for _, inAddr := range inAddrs {
-			addrs, err := resolver.Resolve(ctx, inAddr)
+	if !Matches(in) {
+		return []ma.Multiaddr{in}, nil
+	}
+	var outAddrs []ma.Multiaddr
+	toResolve := []ma.Multiaddr{in}
+
+	for len(toResolve) > 0 {
+		var nextToResolve []ma.Multiaddr
+		for _, a := range toResolve {
+			addrs, err := resolver.Resolve(ctx, a)
 			if err != nil {
 				return nil, err
 			}
-			outAddrs = append(outAddrs, addrs...)
+			for _, addr := range addrs {
+				if Matches(addr) {
+					nextToResolve = append(nextToResolve, addr)
+				} else {
+					outAddrs = append(outAddrs, addr)
+				}
+			}
 		}
+		toResolve = nextToResolve
 	}
 	return outAddrs, nil
 }