AGDNS-4409 Imp upstream tests

Author: makazantsev

dnsproxytest/dnsproxytest.go

@@ -0,0 +1,3 @@
+// Package dnsproxytest provides test utilities and mock implementations
+// for the dnsproxy module interfaces.
+package dnsproxytest

internal/dnsproxytest/interface.go dnsproxytest/interface.gointernal/dnsproxytest/interface.go renamed to dnsproxytest/interface.go

@@ -1,15 +1,17 @@
 package dnsproxytest
 
 import (
-	"github.com/AdguardTeam/dnsproxy/internal/dnsmsg"
+	"context"
+
+	"github.com/AdguardTeam/dnsproxy/proxy"
 	"github.com/AdguardTeam/dnsproxy/upstream"
 	"github.com/AdguardTeam/golibs/testutil"
 	"github.com/miekg/dns"
+	"github.com/quic-go/quic-go"
+	"github.com/quic-go/quic-go/qlogwriter"
 )
 
 // Upstream is a mock [upstream.Upstream] implementation for tests.
-//
-// TODO(e.burkov):  Move to golibs.
 type Upstream struct {
 	OnAddress  func() (addr string)
 	OnExchange func(req *dns.Msg) (resp *dns.Msg, err error)
@@ -34,7 +36,55 @@ func (u *Upstream) Close() (err error) {
 	return u.OnClose()
 }
 
-// MessageConstructor is a mock [dnsmsg.MessageConstructor] implementation for
+// Handler is a mock [proxy.Handler] implementation for tests.
+type Handler struct {
+	OnHandle func(ctx context.Context, p *proxy.Proxy, dctx *proxy.DNSContext) (err error)
+}
+
+// type check
+var _ proxy.Handler = (*Handler)(nil)
+
+// ServeDNS implements the [proxy.Handler] interface for *Handler.
+func (h *Handler) ServeDNS(ctx context.Context, p *proxy.Proxy, dctx *proxy.DNSContext) (err error) {
+	return h.OnHandle(ctx, p, dctx)
+}
+
+// Middleware is a mock [proxy.Middleware] implementation for tests.
+type Middleware struct {
+	OnWrap func(h proxy.Handler) (wrapped proxy.Handler)
+}
+
+// type check
+var _ proxy.Middleware = (*Middleware)(nil)
+
+// Wrap implements the [proxy.Middleware] interface for *Middleware.
+func (m *Middleware) Wrap(h proxy.Handler) (wrapped proxy.Handler) {
+	return m.OnWrap(h)
+}
+
+// QUICTracer is a mock [upstream.QUICTracer] implementation for tests.
+type QUICTracer struct {
+	OnTraceForConnection func(
+		ctx context.Context,
+		isClient bool,
+		connID quic.ConnectionID,
+	) (trace qlogwriter.Trace)
+}
+
+// type check
+var _ upstream.QUICTracer = (*QUICTracer)(nil)
+
+// TraceForConnection implements the [upstream.QUICTracer] interface for
+// *QUICTracer.
+func (t *QUICTracer) TraceForConnection(
+	ctx context.Context,
+	isClient bool,
+	connID quic.ConnectionID,
+) (trace qlogwriter.Trace) {
+	return t.OnTraceForConnection(ctx, isClient, connID)
+}
+
+// MessageConstructor is a mock [proxy.MessageConstructor] implementation for
 // tests.
 type MessageConstructor struct {
 	OnNewMsgNXDOMAIN       func(req *dns.Msg) (resp *dns.Msg)
@@ -44,8 +94,8 @@ type MessageConstructor struct {
 	OnNewMsgFORMERR        func(req *dns.Msg) (resp *dns.Msg)
 }
 
-// NewMessageConstructor creates a new *TestMessageConstructor with all it's
-// methods set to panic.
+// NewMessageConstructor creates a new *MessageConstructor with all its methods
+// set to panic.
 func NewMessageConstructor() (c *MessageConstructor) {
 	return &MessageConstructor{
 		OnNewMsgNXDOMAIN: func(req *dns.Msg) (_ *dns.Msg) {
@@ -67,7 +117,7 @@ func NewMessageConstructor() (c *MessageConstructor) {
 }
 
 // type check
-var _ dnsmsg.MessageConstructor = (*MessageConstructor)(nil)
+var _ proxy.MessageConstructor = (*MessageConstructor)(nil)
 
 // NewMsgNXDOMAIN implements the [proxy.MessageConstructor] interface for
 // *MessageConstructor.

internal/dnsproxytest/dnsproxytest.go

@@ -1,5 +1,5 @@
-// Package dnsproxytest provides a set of test utilities for the dnsproxy
-// module.
+// Package dnsproxytest provides test utilities and mock implementations
+// for the dnsproxy module interfaces.
 package dnsproxytest
 
 import (
@@ -15,6 +15,8 @@ import (
 // that can be used for testing.  Note that there is theoretically a TOCTTOU
 // race here: the port may be reoccupied between the time it is released and the
 // time the caller binds to it.
+//
+// TODO(m.kazantsev):  Move to the top-level dnsproxytest package.
 func NewFreePort(tb testing.TB) (p uint) {
 	tb.Helper()
 

internal/middleware/middleware_internal_test.go

@@ -8,8 +8,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/AdguardTeam/dnsproxy/dnsproxytest"
 	"github.com/AdguardTeam/dnsproxy/internal/dnsmsg"
-	"github.com/AdguardTeam/dnsproxy/internal/dnsproxytest"
 	"github.com/AdguardTeam/golibs/errors"
 	"github.com/AdguardTeam/golibs/logutil/slogutil"
 	"github.com/AdguardTeam/golibs/netutil"

proxy/bogusnxdomain_internal_test.go

@@ -14,11 +14,19 @@ import (
 )
 
 func TestProxy_IsBogusNXDomain(t *testing.T) {
+	var ans []dns.RR
+
+	onExchange := newECSReplyHandler(&ans, nil, nil)
+	u := newTestECSUpstream(onExchange)
+
+	upsConf := newTestUpstreamConfig(t, defaultTimeout, testDefaultUpstreamAddr)
+	upsConf.Upstreams = []upstream.Upstream{u}
+
 	prx := mustNew(t, &Config{
 		Logger:         testLogger,
 		UDPListenAddr:  []*net.UDPAddr{net.UDPAddrFromAddrPort(localhostAnyPort)},
 		TCPListenAddr:  []*net.TCPAddr{net.TCPAddrFromAddrPort(localhostAnyPort)},
-		UpstreamConfig: newTestUpstreamConfig(t, defaultTimeout, testDefaultUpstreamAddr),
+		UpstreamConfig: upsConf,
 		TrustedProxies: defaultTrustedProxies,
 		CacheEnabled:   true,
 		BogusNXDomain: []netip.Prefix{
@@ -77,19 +85,16 @@ func TestProxy_IsBogusNXDomain(t *testing.T) {
 		wantRcode: dns.RcodeSuccess,
 	}}
 
-	u := testUpstream{}
-	prx.upstreamConf.Upstreams = []upstream.Upstream{&u}
-
 	servicetest.RequireRun(t, prx, testTimeout)
 
 	d := &DNSContext{
 		Req: newHostTestMessage("host"),
 	}
 
 	for _, tc := range testCases {
-		u.ans = tc.ans
-
 		t.Run(tc.name, func(t *testing.T) {
+			ans = tc.ans
+
 			err := prx.Resolve(testutil.ContextWithTimeout(t, defaultTimeout), d)
 			require.NoError(t, err)
 			require.NotNil(t, d.Res)

proxy/cache_internal_test.go

@@ -9,7 +9,6 @@ import (
 	"testing"
 	"time"
 
-	"github.com/AdguardTeam/dnsproxy/internal/dnsproxytest"
 	"github.com/AdguardTeam/dnsproxy/upstream"
 	"github.com/AdguardTeam/golibs/netutil"
 	"github.com/AdguardTeam/golibs/testutil"
@@ -28,7 +27,7 @@ const testUpsAddr = "https://upstream.address"
 
 // upstreamWithAddr is a [dnsproxytest.Upstream] that is only expected to be
 // used to get its address.
-var upstreamWithAddr = &dnsproxytest.Upstream{
+var upstreamWithAddr = &testUpstream{
 	OnExchange: func(m *dns.Msg) (_ *dns.Msg, _ error) { panic(testutil.UnexpectedCall(m)) },
 	OnClose:    func() (_ error) { panic(testutil.UnexpectedCall()) },
 	OnAddress:  func() (addr string) { return testUpsAddr },
@@ -361,14 +360,17 @@ func TestCacheExpiration(t *testing.T) {
 }
 
 func TestCacheExpirationWithTTLOverride(t *testing.T) {
-	u := testUpstream{}
+	var ans []dns.RR
+
+	onExchange := newECSReplyHandler(&ans, nil, nil)
+	u := newTestECSUpstream(onExchange)
 
 	dnsProxy := mustNew(t, &Config{
 		Logger:        testLogger,
 		UDPListenAddr: []*net.UDPAddr{net.UDPAddrFromAddrPort(localhostAnyPort)},
 		TCPListenAddr: []*net.TCPAddr{net.TCPAddrFromAddrPort(localhostAnyPort)},
 		UpstreamConfig: &UpstreamConfig{
-			Upstreams: []upstream.Upstream{&u},
+			Upstreams: []upstream.Upstream{u},
 		},
 		TrustedProxies: defaultTrustedProxies,
 		CacheEnabled:   true,
@@ -385,7 +387,7 @@ func TestCacheExpirationWithTTLOverride(t *testing.T) {
 		d.Req = newHostTestMessage("host")
 		d.Addr = netip.AddrPort{}
 
-		u.ans = []dns.RR{&dns.A{
+		ans = []dns.RR{&dns.A{
 			Hdr: dns.RR_Header{
 				Rrtype: dns.TypeA,
 				Name:   "host.",
@@ -409,7 +411,7 @@ func TestCacheExpirationWithTTLOverride(t *testing.T) {
 		d.Req = newHostTestMessage("host2")
 		d.Addr = netip.AddrPort{}
 
-		u.ans = []dns.RR{&dns.A{
+		ans = []dns.RR{&dns.A{
 			Hdr: dns.RR_Header{
 				Rrtype: dns.TypeA,
 				Name:   "host2.",

proxy/dns64_internal_test.go

@@ -6,7 +6,6 @@ import (
 	"sync"
 	"testing"
 
-	"github.com/AdguardTeam/dnsproxy/internal/dnsproxytest"
 	"github.com/AdguardTeam/dnsproxy/upstream"
 	"github.com/AdguardTeam/golibs/netutil"
 	"github.com/AdguardTeam/golibs/testutil"
@@ -20,7 +19,7 @@ const ipv4OnlyFqdn = "ipv4.only."
 
 func TestDNS64Race(t *testing.T) {
 	ans := newRR(t, ipv4OnlyFqdn, dns.TypeA, 3600, net.ParseIP("1.2.3.4"))
-	ups := &dnsproxytest.Upstream{
+	ups := &testUpstream{
 		OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
 			resp = (&dns.Msg{}).SetReply(req)
 			if req.Question[0].Qtype == dns.TypeA {
@@ -32,7 +31,7 @@ func TestDNS64Race(t *testing.T) {
 		OnAddress: func() (addr string) { return "fake.address" },
 		OnClose:   func() (err error) { return nil },
 	}
-	localUps := &dnsproxytest.Upstream{
+	localUps := &testUpstream{
 		OnExchange: func(m *dns.Msg) (_ *dns.Msg, _ error) { panic(testutil.UnexpectedCall(m)) },
 		OnAddress:  func() (addr string) { return "fake.address" },
 		OnClose:    func() (err error) { return nil },
@@ -180,7 +179,7 @@ func TestProxy_Resolve_dns64(t *testing.T) {
 
 	pt := testutil.PanicT{}
 	newUps := func(answers answerMap) (u upstream.Upstream) {
-		return &dnsproxytest.Upstream{
+		return &testUpstream{
 			OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
 				q := req.Question[0]
 				require.Contains(pt, answers, q.Qtype)
@@ -200,7 +199,7 @@ func TestProxy_Resolve_dns64(t *testing.T) {
 	}
 
 	localRR := newRR(t, ptr64Domain, dns.TypePTR, 3600, domainPointed)
-	localUps := &dnsproxytest.Upstream{
+	localUps := &testUpstream{
 		OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
 			require.Equal(pt, req.Question[0].Name, ptr64Domain)
 			resp = (&dns.Msg{}).SetReply(req)

proxy/exchange_internal_test.go

@@ -8,7 +8,6 @@ import (
 	"testing"
 	"time"
 
-	"github.com/AdguardTeam/dnsproxy/internal/dnsproxytest"
 	"github.com/AdguardTeam/dnsproxy/upstream"
 	"github.com/AdguardTeam/golibs/netutil"
 	"github.com/AdguardTeam/golibs/testutil"
@@ -24,7 +23,7 @@ import (
 func newUpstreamWithErrorRate(rate uint, name string) (u upstream.Upstream) {
 	var n uint
 
-	return &dnsproxytest.Upstream{
+	return &testUpstream{
 		OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
 			n++
 			if n%rate == 0 {
@@ -87,7 +86,7 @@ func TestProxy_Exchange_loadBalance(t *testing.T) {
 		},
 	}
 
-	fastUps := &dnsproxytest.Upstream{
+	fastUps := &testUpstream{
 		OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
 			currentNow = zeroTime.Add(testRTT / 100)
 
@@ -96,7 +95,7 @@ func TestProxy_Exchange_loadBalance(t *testing.T) {
 		OnAddress: func() (addr string) { return "fast" },
 		OnClose:   func() (_ error) { panic(testutil.UnexpectedCall()) },
 	}
-	slowerUps := &dnsproxytest.Upstream{
+	slowerUps := &testUpstream{
 		OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
 			currentNow = zeroTime.Add(testRTT / 10)
 
@@ -105,7 +104,7 @@ func TestProxy_Exchange_loadBalance(t *testing.T) {
 		OnAddress: func() (addr string) { return "slower" },
 		OnClose:   func() (_ error) { panic(testutil.UnexpectedCall()) },
 	}
-	slowestUps := &dnsproxytest.Upstream{
+	slowestUps := &testUpstream{
 		OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
 			currentNow = zeroTime.Add(testRTT / 2)
 
@@ -115,20 +114,20 @@ func TestProxy_Exchange_loadBalance(t *testing.T) {
 		OnClose:   func() (_ error) { panic(testutil.UnexpectedCall()) },
 	}
 
-	err1Ups := &dnsproxytest.Upstream{
+	err1Ups := &testUpstream{
 		OnExchange: func(_ *dns.Msg) (r *dns.Msg, err error) { return nil, assert.AnError },
 		OnAddress:  func() (addr string) { return "error1" },
 		OnClose:    func() (_ error) { panic(testutil.UnexpectedCall()) },
 	}
-	err2Ups := &dnsproxytest.Upstream{
+	err2Ups := &testUpstream{
 		OnExchange: func(_ *dns.Msg) (r *dns.Msg, err error) { return nil, assert.AnError },
 		OnAddress:  func() (addr string) { return "error2" },
 		OnClose:    func() (_ error) { panic(testutil.UnexpectedCall()) },
 	}
 
 	singleError := &sync.Once{}
 	// fastestUps responds with an error on the first request.
-	fastestUps := &dnsproxytest.Upstream{
+	fastestUps := &testUpstream{
 		OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
 			singleError.Do(func() { err = assert.AnError })
 			currentNow = zeroTime.Add(testRTT / 200)

proxy/handler_internal_test.go

@@ -17,7 +17,7 @@ func TestFilteringHandler(t *testing.T) {
 	m := &sync.RWMutex{}
 	blockResponse := false
 
-	reqHandler := &TestHandler{
+	reqHandler := &testHandler{
 		OnHandle: func(ctx context.Context, p *Proxy, d *DNSContext) (err error) {
 			m.Lock()
 			defer m.Unlock()

proxy/pending_test.go

@@ -8,7 +8,7 @@ import (
 	"testing"
 	"time"
 
-	"github.com/AdguardTeam/dnsproxy/internal/dnsproxytest"
+	"github.com/AdguardTeam/dnsproxy/dnsproxytest"
 	"github.com/AdguardTeam/dnsproxy/proxy"
 	"github.com/AdguardTeam/dnsproxy/upstream"
 	"github.com/AdguardTeam/golibs/netutil"
@@ -79,7 +79,7 @@ func TestPendingRequests(t *testing.T) {
 	workloadWG := &sync.WaitGroup{}
 	workloadWG.Add(reqsNum)
 
-	reqHandler := &proxy.TestHandler{
+	reqHandler := &dnsproxytest.Handler{
 		OnHandle: func(ctx context.Context, p *proxy.Proxy, d *proxy.DNSContext) (err error) {
 			workloadWG.Done()