dnsproxy/proxy/stats_test.go

package proxy_test

import (
	"net"
	"testing"

	"github.com/AdguardTeam/dnsproxy/dnsproxytest"
	proxytest "github.com/AdguardTeam/dnsproxy/internal/dnsproxytest"
	"github.com/AdguardTeam/dnsproxy/proxy"
	"github.com/AdguardTeam/dnsproxy/upstream"
	"github.com/AdguardTeam/golibs/errors"
	"github.com/AdguardTeam/golibs/testutil"
	"github.com/miekg/dns"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestDNSContext_QueryStatistics(t *testing.T) {
	t.Parallel()

	testReq := &dns.Msg{
		Question: []dns.Question{{
			Name:   "test.",
			Qtype:  dns.TypeA,
			Qclass: dns.ClassINET,
		}},
	}

	ups := &dnsproxytest.Upstream{
		OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
			return (&dns.Msg{}).SetReply(req), nil
		},
		OnAddress: func() (addr string) { return "upstream" },
		OnClose:   func() (err error) { return nil },
	}

	failUps := &dnsproxytest.Upstream{
		OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
			return nil, errors.Error("exchange error")
		},
		OnAddress: func() (addr string) { return "fail.upstream" },
		OnClose:   func() (err error) { return nil },
	}

	conf := &proxy.Config{
		Logger:         testLogger,
		UDPListenAddr:  []*net.UDPAddr{net.UDPAddrFromAddrPort(proxytest.LocalhostAnyPort)},
		TCPListenAddr:  []*net.TCPAddr{net.TCPAddrFromAddrPort(proxytest.LocalhostAnyPort)},
		TrustedProxies: proxytest.DefaultTrustedProxies,
	}

	testCases := []struct {
		wantErr           assert.ErrorAssertionFunc
		wantMainErr       assert.BoolAssertionFunc
		wantFallbackErr   assert.BoolAssertionFunc
		config            *proxy.UpstreamConfig
		fallbackConfig    *proxy.UpstreamConfig
		name              string
		mode              proxy.UpstreamMode
		wantMainCount     int
		wantFallbackCount int
	}{{
		wantErr:         assert.NoError,
		wantMainErr:     assert.False,
		wantFallbackErr: assert.False,
		config: &proxy.UpstreamConfig{
			Upstreams: []upstream.Upstream{ups},
		},
		fallbackConfig: &proxy.UpstreamConfig{
			Upstreams: []upstream.Upstream{ups},
		},
		name:              "load_balance_success",
		mode:              proxy.UpstreamModeLoadBalance,
		wantMainCount:     1,
		wantFallbackCount: 0,
	}, {
		wantErr:         assert.Error,
		wantMainErr:     assert.True,
		wantFallbackErr: assert.True,
		config: &proxy.UpstreamConfig{
			Upstreams: []upstream.Upstream{failUps},
		},
		fallbackConfig: &proxy.UpstreamConfig{
			Upstreams: []upstream.Upstream{failUps, failUps},
		},
		name:              "load_balance_bad",
		mode:              proxy.UpstreamModeLoadBalance,
		wantMainCount:     1,
		wantFallbackCount: 2,
	}, {
		wantErr:         assert.NoError,
		wantMainErr:     assert.False,
		wantFallbackErr: assert.False,
		config: &proxy.UpstreamConfig{
			Upstreams: []upstream.Upstream{ups, failUps},
		},
		fallbackConfig: &proxy.UpstreamConfig{
			Upstreams: []upstream.Upstream{ups},
		},
		name:              "parallel_success",
		mode:              proxy.UpstreamModeParallel,
		wantMainCount:     1,
		wantFallbackCount: 0,
	}, {
		wantErr:         assert.NoError,
		wantMainErr:     assert.True,
		wantFallbackErr: assert.False,
		config: &proxy.UpstreamConfig{
			Upstreams: []upstream.Upstream{failUps},
		},
		fallbackConfig: &proxy.UpstreamConfig{
			Upstreams: []upstream.Upstream{ups},
		},
		name:              "parallel_bad_fallback_success",
		mode:              proxy.UpstreamModeParallel,
		wantMainCount:     1,
		wantFallbackCount: 1,
	}, {
		wantErr:         assert.Error,
		wantMainErr:     assert.True,
		wantFallbackErr: assert.True,
		config: &proxy.UpstreamConfig{
			Upstreams: []upstream.Upstream{failUps, failUps},
		},
		fallbackConfig: &proxy.UpstreamConfig{
			Upstreams: []upstream.Upstream{failUps, failUps, failUps},
		},
		name:              "parallel_bad",
		mode:              proxy.UpstreamModeParallel,
		wantMainCount:     2,
		wantFallbackCount: 3,
	}, {
		wantErr:         assert.NoError,
		wantMainErr:     assert.False,
		wantFallbackErr: assert.False,
		config: &proxy.UpstreamConfig{
			Upstreams: []upstream.Upstream{ups},
		},
		fallbackConfig: &proxy.UpstreamConfig{
			Upstreams: []upstream.Upstream{ups},
		},
		name:              "fastest_single_success",
		mode:              proxy.UpstreamModeFastestAddr,
		wantMainCount:     1,
		wantFallbackCount: 0,
	}, {
		wantErr:         assert.NoError,
		wantMainErr:     assert.False,
		wantFallbackErr: assert.False,
		config: &proxy.UpstreamConfig{
			Upstreams: []upstream.Upstream{ups, ups},
		},
		fallbackConfig: &proxy.UpstreamConfig{
			Upstreams: []upstream.Upstream{ups},
		},
		name:              "fastest_multiple_success",
		mode:              proxy.UpstreamModeFastestAddr,
		wantMainCount:     2,
		wantFallbackCount: 0,
	}, {
		wantErr:         assert.NoError,
		wantMainErr:     assert.True,
		wantFallbackErr: assert.False,
		config: &proxy.UpstreamConfig{
			Upstreams: []upstream.Upstream{ups, failUps},
		},
		fallbackConfig: &proxy.UpstreamConfig{
			Upstreams: []upstream.Upstream{ups},
		},
		name:              "fastest_mixed_success",
		mode:              proxy.UpstreamModeFastestAddr,
		wantMainCount:     2,
		wantFallbackCount: 0,
	}, {
		wantErr:         assert.Error,
		wantMainErr:     assert.True,
		wantFallbackErr: assert.True,
		config: &proxy.UpstreamConfig{
			Upstreams: []upstream.Upstream{failUps, failUps},
		},
		fallbackConfig: &proxy.UpstreamConfig{
			Upstreams: []upstream.Upstream{failUps, failUps, failUps},
		},
		name:              "fastest_multiple_bad",
		mode:              proxy.UpstreamModeFastestAddr,
		wantMainCount:     2,
		wantFallbackCount: 3,
	}, {
		wantErr:         assert.NoError,
		wantMainErr:     assert.True,
		wantFallbackErr: assert.False,
		config: &proxy.UpstreamConfig{
			Upstreams: []upstream.Upstream{failUps, failUps},
		},
		fallbackConfig: &proxy.UpstreamConfig{
			Upstreams: []upstream.Upstream{ups},
		},
		name:              "fastest_bad_fallback_success",
		mode:              proxy.UpstreamModeFastestAddr,
		wantMainCount:     2,
		wantFallbackCount: 1,
	}}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			conf.UpstreamConfig = tc.config
			conf.Fallbacks = tc.fallbackConfig
			conf.UpstreamMode = tc.mode

			p, err := proxy.New(conf)
			require.NoError(t, err)

			d := &proxy.DNSContext{Req: testReq}

			err = p.Resolve(testutil.ContextWithTimeout(t, defaultTimeout), d)
			tc.wantErr(t, err)

			stats := d.QueryStatistics()
			assertQueryStats(
				t,
				stats,
				tc.wantMainCount,
				tc.wantMainErr,
				tc.wantFallbackCount,
				tc.wantFallbackErr,
			)
		})
	}
}

// assertQueryStats asserts the statistics using the provided parameters.
func assertQueryStats(
	tb testing.TB,
	stats *proxy.QueryStatistics,
	wantMainCount int,
	wantMainErr assert.BoolAssertionFunc,
	wantFallbackCount int,
	wantFallbackErr assert.BoolAssertionFunc,
) {
	tb.Helper()

	main := stats.Main()
	assert.Lenf(tb, main, wantMainCount, "main stats count")

	fallback := stats.Fallback()
	assert.Lenf(tb, fallback, wantFallbackCount, "fallback stats count")

	wantMainErr(tb, isErrorInStats(main), "main err")
	wantFallbackErr(tb, isErrorInStats(fallback), "fallback err")
}

// isErrorInStats is a helper function for tests that returns true if the
// upstream statistics contain an DNS lookup error.
func isErrorInStats(stats []*proxy.UpstreamStatistics) (ok bool) {
	for _, u := range stats {
		if u.Error != nil {
			return true
		}
	}

	return false
}