dnsforward: fix safesearch test

This commit is contained in:
Stanislav Chzhen 2024-11-22 19:36:27 +03:00
parent abb738013a
commit efabbe318e
1 changed files with 18 additions and 11 deletions

View File

@ -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{}
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)