dnsforward: imp tests

annguyen0 · Apr 23, 2024 · af2507b · af2507b
1 parent 67fc9d3
commit af2507b
Showing 1 changed file with 27 additions and 1 deletion.
diff --git a/internal/dnsforward/beforerequest_internal_test.go b/internal/dnsforward/beforerequest_internal_test.go
@@ -4,7 +4,9 @@ import (
 	"crypto/tls"
 	"net"
 	"testing"
+	"time"
 
+	"github.com/AdguardTeam/AdGuardHome/internal/aghtest"
 	"github.com/AdguardTeam/AdGuardHome/internal/filtering"
 	"github.com/AdguardTeam/dnsproxy/proxy"
 	"github.com/miekg/dns"
@@ -104,6 +106,15 @@ func TestServer_HandleBefore(t *testing.T) {
 		wantRCode:         dns.RcodeRefused,
 	}}
 
+	localRR := newRR(t, testFQDN, dns.TypeA, 3600, net.IP{1, 2, 3, 4})
+	localUpsHdlr := dns.HandlerFunc(func(w dns.ResponseWriter, req *dns.Msg) {
+		resp := (&dns.Msg{}).SetReply(req)
+		resp.Answer = []dns.RR{localRR}
+
+		require.NoError(t, w.WriteMsg(resp))
+	})
+	localUpsAddr := aghtest.StartLocalhostUpstream(t, localUpsHdlr).String()
+
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
@@ -113,6 +124,8 @@ func TestServer_HandleBefore(t *testing.T) {
 				ServerName:     tlsServerName,
 			})
 
+			s.conf.UpstreamDNS = []string{localUpsAddr}
+
 			s.conf.AllowedClients = tc.allowedClients
 			s.conf.DisallowedClients = tc.disallowedClients
 			s.conf.BlockedHosts = tc.blockedHosts
@@ -130,6 +143,7 @@ func TestServer_HandleBefore(t *testing.T) {
 			client := &dns.Client{
 				Net:       "tcp-tls",
 				TLSConfig: tlsConfig,
+				Timeout:   200 * time.Millisecond,
 			}
 
 			req := createTestMessage(tc.host)
@@ -215,6 +229,15 @@ func TestServer_HandleBefore_udp(t *testing.T) {
 		wantTimeout:       true,
 	}}
 
+	localRR := newRR(t, testFQDN, dns.TypeA, 3600, net.IP{1, 2, 3, 4})
+	localUpsHdlr := dns.HandlerFunc(func(w dns.ResponseWriter, req *dns.Msg) {
+		resp := (&dns.Msg{}).SetReply(req)
+		resp.Answer = []dns.RR{localRR}
+
+		require.NoError(t, w.WriteMsg(resp))
+	})
+	localUpsAddr := aghtest.StartLocalhostUpstream(t, localUpsHdlr).String()
+
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
@@ -225,6 +248,7 @@ func TestServer_HandleBefore_udp(t *testing.T) {
 				UDPListenAddrs: []*net.UDPAddr{{}},
 				TCPListenAddrs: []*net.TCPAddr{{}},
 				Config: Config{
+					UpstreamDNS:      []string{localUpsAddr},
 					UpstreamMode:     UpstreamModeLoadBalance,
 					EDNSClientSubnet: &EDNSClientSubnet{Enabled: false},
 				},
@@ -241,7 +265,8 @@ func TestServer_HandleBefore_udp(t *testing.T) {
 			startDeferStop(t, s)
 
 			client := &dns.Client{
-				Net: "udp",
+				Net:     "udp",
+				Timeout: 200 * time.Millisecond,
 			}
 
 			req := createTestMessage(tc.host)
@@ -255,6 +280,7 @@ func TestServer_HandleBefore_udp(t *testing.T) {
 
 				assert.Nil(t, reply)
 			} else {
+				require.NoError(t, err)
 				require.NotNil(t, reply)
 
 				assert.Equal(t, dns.RcodeSuccess, reply.Rcode)