resolver/mock_test.go at main · nsmithuk/resolver · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
package resolver

import (
	"context"
	"github.com/miekg/dns"
	"github.com/stretchr/testify/mock"
)

// Mock expiringExchanger to simulate pool expiration behavior and DNS message exchange
type MockExpiringExchanger struct {
	mock.Mock
}

// Mock the expired function
func (m *MockExpiringExchanger) expired() bool {
	args := m.Called()
	return args.Bool(0)
}

// Mock the exchange function to simulate DNS message exchange
func (m *MockExpiringExchanger) exchange(ctx context.Context, msg *dns.Msg) *Response {
	args := m.Called(ctx, msg)
	return args.Get(0).(*Response)
}

//--------------------------------------------------------------------------

type mockZoneStore struct {
	mockAdd      func(z zone)
	mockGet      func(name string) zone
	mockCount    func() int
	mockZoneList func(name string) []zone
}

func (m mockZoneStore) getZoneList(name string) []zone {
	return m.mockZoneList(name)
}
func (m mockZoneStore) get(name string) zone {
	return m.mockGet(name)
}
func (m mockZoneStore) add(z zone) {
	m.mockAdd(z)
}
func (m mockZoneStore) count() int {
	return m.mockCount()
}

//--------------------------------------------------------------------------

type mockExchanger struct {
	mockExchange func(context.Context, *dns.Msg) *Response
}

func (m *mockExchanger) exchange(ctx context.Context, qmsg *dns.Msg) *Response {
	return m.mockExchange(ctx, qmsg)
}

//--------------------------------------------------------------------------

type mockZone struct {
	mockName     func() string
	mockParent   func() string
	mockExpired  func() bool
	mockClone    func(name, parent string) zone
	mockSoa      func(ctx context.Context, name string) (*dns.SOA, error)
	mockDnskeys  func(ctx context.Context) ([]dns.RR, error)
	mockExchange func(ctx context.Context, m *dns.Msg) *Response
}

func (z *mockZone) name() string {
	return z.mockName()
}

func (z *mockZone) parent() string {
	return z.mockParent()
}

func (z *mockZone) expired() bool {
	return z.mockExpired()
}

func (z *mockZone) clone(name, parent string) zone {
	return z.mockClone(name, parent)
}

func (z *mockZone) soa(ctx context.Context, name string) (*dns.SOA, error) {
	return z.mockSoa(ctx, name)
}

func (z *mockZone) dnskeys(ctx context.Context) ([]dns.RR, error) {
	return z.mockDnskeys(ctx)
}

func (z *mockZone) exchange(ctx context.Context, m *dns.Msg) *Response {
	return z.mockExchange(ctx, m)
}