From 006391c9c089a2d9fe7ad7157e609898be2ee225 Mon Sep 17 00:00:00 2001
From: Ainar Garipov <A.Garipov@AdGuard.COM>
Date: Thu, 3 Oct 2024 14:15:10 +0300
Subject: [PATCH] dnsforward: imp test more

---
 internal/dnsforward/dnsforward_test.go | 18 +++++++-----------
 1 file changed, 7 insertions(+), 11 deletions(-)

diff --git a/internal/dnsforward/dnsforward_test.go b/internal/dnsforward/dnsforward_test.go
index 341c29fe03f..25565503caf 100644
--- a/internal/dnsforward/dnsforward_test.go
+++ b/internal/dnsforward/dnsforward_test.go
@@ -583,19 +583,15 @@ func TestSafeSearch(t *testing.T) {
 		t.Run(tc.host, func(t *testing.T) {
 			req := createTestMessage(tc.host)
 
+			// TODO(a.garipov):  Create our own helper for this.
 			var reply *dns.Msg
-
-			// Use a mutex, since require.Eventually uses goroutines.
-			mu := &sync.Mutex{}
-			require.Eventually(t, func() (ok bool) {
-				mu.Lock()
-				defer mu.Unlock()
-
-				reply, _, err = client.Exchange(req, addr)
-
-				return err == nil
+			once := &sync.Once{}
+			require.EventuallyWithT(t, func(c *assert.CollectT) {
+				r, _, errExch := client.Exchange(req, addr)
+				if assert.NoError(c, errExch) {
+					once.Do(func() { reply = r })
+				}
 			}, testTimeout*10, testTimeout)
-			require.NoErrorf(t, err, "couldn't talk to server %s: %s", addr, err)
 
 			if tc.wantCNAME != "" {
 				require.Len(t, reply.Answer, 2)