dnsforward: fix safesearch test

internal/dnsforward/dnsforward_test.go

@@ -536,10 +536,24 @@ 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) {
+		resp = new(dns.Msg).SetReply(req)
+		resp.Answer = []dns.RR{&dns.A{
+			Hdr: dns.RR_Header{
+				Name:   req.Question[0].Name,
+				Rrtype: dns.TypeA,
+				Class:  dns.ClassINET,
+				Ttl:    60,
+			},
+			A: []byte{1, 2, 3, 4},
+		}}
 
+		return resp, 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})
 
@@ -587,15 +601,8 @@ func TestSafeSearch(t *testing.T) {
 
 			// TODO(a.garipov):  Create our own helper for this.
 			var reply *dns.Msg
-			once := &sync.Once{}
+			reply, err = dns.Exchange(req, addr)
-			require.EventuallyWithT(t, func(c *assert.CollectT) {
+			require.NoError(t, err)
-				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)
 
 			if tc.wantCNAME != "" {
 				require.Len(t, reply.Answer, 2)