all: imp code

AdguardTeam · May 31, 2024 · 9a452ed · 9a452ed
1 parent 9d11c99
commit 9a452ed
Show file tree

Hide file tree

Showing 5 changed files with 38 additions and 13 deletions.
diff --git a/proxy/cache_test.go b/proxy/cache_test.go
@@ -59,7 +59,10 @@ func TestServeCached(t *testing.T) {
 
 	// Create a DNS-over-UDP client connection.
 	addr := dnsProxy.Addr(ProtoUDP)
-	client := &dns.Client{Net: "udp", Timeout: testTimeout}
+	client := &dns.Client{
+		Net:     string(ProtoUDP),
+		Timeout: testTimeout,
+	}
 
 	// Create a DNS request.
 	request := (&dns.Msg{}).SetQuestion("google.com.", dns.TypeA)
@@ -270,6 +273,14 @@ func TestCache_concurrent(t *testing.T) {
 	g.Wait()
 }
 
+const (
+	// cacheTick is a cache check period.
+	cacheTick = 50 * time.Millisecond
+
+	// cacheTimeout is the timeout of cache check.
+	cacheTimeout = 20 * cacheTick
+)
+
 func TestCacheExpiration(t *testing.T) {
 	t.Parallel()
 
@@ -324,7 +335,7 @@ func TestCacheExpiration(t *testing.T) {
 		}
 
 		return true
-	}, 1100*time.Millisecond, 50*time.Millisecond)
+	}, cacheTimeout, cacheTick)
 }
 
 func TestCacheExpirationWithTTLOverride(t *testing.T) {
@@ -609,7 +620,7 @@ func setAndGetCache(t *testing.T, c *cache, g *sync.WaitGroup, host, ip string)
 		ci, _, _ := c.get(dnsMsg)
 
 		return ci == nil
-	}, 1100*time.Millisecond, 100*time.Millisecond, "cache for %s should already be removed", host)
+	}, cacheTimeout, cacheTick, "cache for %s should already be removed", host)
 }
 
 func TestCache_getWithSubnet(t *testing.T) {

diff --git a/proxy/handler_test.go b/proxy/handler_test.go
@@ -52,7 +52,10 @@ func TestFilteringHandler(t *testing.T) {
 
 	// Create a DNS-over-UDP client connection
 	addr := dnsProxy.Addr(ProtoUDP)
-	client := &dns.Client{Net: "udp", Timeout: testTimeout}
+	client := &dns.Client{
+		Net:     string(ProtoUDP),
+		Timeout: testTimeout,
+	}
 
 	// Send the first message (not blocked)
 	req := newTestMessage()

diff --git a/proxy/proxy_test.go b/proxy/proxy_test.go
@@ -901,7 +901,10 @@ func TestRefuseAny(t *testing.T) {
 
 	// Create a DNS-over-UDP client connection
 	addr := dnsProxy.Addr(ProtoUDP)
-	client := &dns.Client{Net: "udp", Timeout: testTimeout}
+	client := &dns.Client{
+		Net:     string(ProtoUDP),
+		Timeout: testTimeout,
+	}
 
 	// Create a DNS request
 	request := (&dns.Msg{
@@ -935,7 +938,10 @@ func TestInvalidDNSRequest(t *testing.T) {
 	testutil.CleanupAndRequireSuccess(t, func() (err error) { return dnsProxy.Shutdown(ctx) })
 
 	// Create a DNS-over-UDP client connection
-	client := &dns.Client{Net: "udp", Timeout: testTimeout}
+	client := &dns.Client{
+		Net:     string(ProtoUDP),
+		Timeout: testTimeout,
+	}
 
 	// Create a DNS request
 	request := &dns.Msg{
@@ -955,7 +961,10 @@ func TestResponseInRequest(t *testing.T) {
 	dnsProxy := mustStartDefaultProxy(t)
 
 	addr := dnsProxy.Addr(ProtoUDP)
-	client := &dns.Client{Net: "udp", Timeout: testTimeout}
+	client := &dns.Client{
+		Net:     string(ProtoUDP),
+		Timeout: testTimeout,
+	}
 
 	req := newTestMessage()
 	req.Response = true
@@ -973,7 +982,10 @@ func TestNoQuestion(t *testing.T) {
 	dnsProxy := mustStartDefaultProxy(t)
 
 	addr := dnsProxy.Addr(ProtoUDP)
-	client := &dns.Client{Net: "udp", Timeout: testTimeout}
+	client := &dns.Client{
+		Net:     string(ProtoUDP),
+		Timeout: testTimeout,
+	}
 
 	req := newTestMessage()
 	req.Question = nil

diff --git a/proxy/ratelimit_test.go b/proxy/ratelimit_test.go
@@ -30,7 +30,10 @@ func TestRatelimitingProxy(t *testing.T) {
 
 	// Create a DNS-over-UDP client connection
 	addr := dnsProxy.Addr(ProtoUDP)
-	client := &dns.Client{Net: "udp", Timeout: testTimeout}
+	client := &dns.Client{
+		Net:     string(ProtoUDP),
+		Timeout: testTimeout,
+	}
 
 	// Send the first message (not blocked)
 	req := newTestMessage()

diff --git a/upstream/dnscrypt_internal_test.go b/upstream/dnscrypt_internal_test.go
@@ -19,8 +19,6 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-// Helpers
-
 // dnsCryptHandlerFunc is a function-based implementation of the
 // [dnscrypt.Handler] interface.
 type dnsCryptHandlerFunc func(w dnscrypt.ResponseWriter, r *dns.Msg) (err error)
@@ -89,8 +87,6 @@ func startTestDNSCryptServer(
 	return stamp
 }
 
-// Tests
-
 func TestUpstreamDNSCrypt(t *testing.T) {
 	t.Parallel()