From d4ca14806e708cf6703d8c69e218dd80bdd43d76 Mon Sep 17 00:00:00 2001
From: Stanislav Chzhen <s.chzhen@adguard.com>
Date: Mon, 25 Nov 2024 14:31:21 +0300
Subject: [PATCH] Pull request 2309: fix-safesearch-test

Squashed commit of the following:

commit a5b6c831d86ee5264c0918bf12ddef7c78f78f38
Author: Stanislav Chzhen <s.chzhen@adguard.com>
Date:   Mon Nov 25 13:45:17 2024 +0300

    dnsforward: imp tests

commit 1ca353423ca4dcc7616475fbf12c9b5b9450a5ea
Author: Stanislav Chzhen <s.chzhen@adguard.com>
Date:   Fri Nov 22 20:08:28 2024 +0300

    dnsforward: imp tests

commit 82851aa825818159b5b1583e2716d6656d633b03
Author: Stanislav Chzhen <s.chzhen@adguard.com>
Date:   Fri Nov 22 19:54:53 2024 +0300

    dnsforward: imp tests

commit efabbe318e7520eba65b595f11e2fb8effaeec13
Author: Stanislav Chzhen <s.chzhen@adguard.com>
Date:   Fri Nov 22 19:36:27 2024 +0300

    dnsforward: fix safesearch test
---
 internal/dnsforward/dnsforward_test.go | 27 ++++++++++++++------------
 1 file changed, 15 insertions(+), 12 deletions(-)

diff --git a/internal/dnsforward/dnsforward_test.go b/internal/dnsforward/dnsforward_test.go
index 548c79da..6227dd09 100644
--- a/internal/dnsforward/dnsforward_test.go
+++ b/internal/dnsforward/dnsforward_test.go
@@ -500,6 +500,10 @@ func TestServerRace(t *testing.T) {
 }
 
 func TestSafeSearch(t *testing.T) {
+	const (
+		googleSafeSearch = "forcesafesearch.google.com."
+	)
+
 	safeSearchConf := filtering.SafeSearchConfig{
 		Enabled: true,
 		Google:  true,
@@ -536,10 +540,17 @@ func TestSafeSearch(t *testing.T) {
 		ServePlainDNS: true,
 	}
 	s := createTestServer(t, filterConf, forwardConf)
-	startDeferStop(t, s)
 
+	ups := aghtest.NewUpstreamMock(func(req *dns.Msg) (resp *dns.Msg, err error) {
+		pt := testutil.PanicT{}
+		assert.Equal(pt, googleSafeSearch, req.Question[0].Name)
+
+		return aghtest.MatchedResponse(req, dns.TypeA, googleSafeSearch, "1.2.3.4"), nil
+	})
+	s.conf.UpstreamConfig.Upstreams = []upstream.Upstream{ups}
+
+	startDeferStop(t, s)
 	addr := s.dnsProxy.Addr(proxy.ProtoUDP).String()
-	client := &dns.Client{}
 
 	yandexIP := netip.AddrFrom4([4]byte{213, 180, 193, 56})
 
@@ -585,17 +596,9 @@ 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
-			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 })
-				} else {
-					t.Logf("got error: %v", errExch)
-				}
-			}, testTimeout*10, testTimeout)
+			reply, err = dns.Exchange(req, addr)
+			require.NoError(t, err)
 
 			if tc.wantCNAME != "" {
 				require.Len(t, reply.Answer, 2)