68 lines
1.7 KiB
Go
68 lines
1.7 KiB
Go
|
package dns
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"net"
|
||
|
|
"sync/atomic"
|
||
|
|
"testing"
|
||
|
|
)
|
||
|
|
|
||
|
|
type fakeLookup struct {
|
||
|
|
calls atomic.Int32
|
||
|
|
cnameErr error
|
||
|
|
ipErr error
|
||
|
|
}
|
||
|
|
|
||
|
|
func (f *fakeLookup) LookupCNAME(ctx context.Context, host string) (string, error) {
|
||
|
|
f.calls.Add(1)
|
||
|
|
if f.cnameErr != nil {
|
||
|
|
return "", f.cnameErr
|
||
|
|
}
|
||
|
|
return host, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (f *fakeLookup) LookupIP(ctx context.Context, host string) ([]net.IP, error) {
|
||
|
|
f.calls.Add(1)
|
||
|
|
if f.ipErr != nil {
|
||
|
|
return nil, f.ipErr
|
||
|
|
}
|
||
|
|
return []net.IP{net.ParseIP("1.2.3.4")}, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestServerPoolRoundRobin(t *testing.T) {
|
||
|
|
a, b := &fakeLookup{}, &fakeLookup{}
|
||
|
|
p := &ServerPool{resolvers: []Resolver{a, b}}
|
||
|
|
for i := 0; i < 4; i++ {
|
||
|
|
if _, err := p.LookupIP(context.Background(), "example.com"); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if a.calls.Load() != 2 || b.calls.Load() != 2 {
|
||
|
|
t.Fatalf("round robin calls: a=%d b=%d, want 2/2", a.calls.Load(), b.calls.Load())
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestServerPoolFailover(t *testing.T) {
|
||
|
|
a := &fakeLookup{ipErr: &net.DNSError{Err: "timeout", IsTimeout: true}}
|
||
|
|
b := &fakeLookup{}
|
||
|
|
p := &ServerPool{resolvers: []Resolver{a, b}}
|
||
|
|
ips, err := p.LookupIP(context.Background(), "example.com")
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if len(ips) != 1 || a.calls.Load() != 1 {
|
||
|
|
t.Fatalf("expected failover to second server, a.calls=%d", a.calls.Load())
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestServerPoolNoFailoverOnNXDOMAIN(t *testing.T) {
|
||
|
|
a := &fakeLookup{ipErr: &net.DNSError{Err: "no such host", IsNotFound: true}}
|
||
|
|
b := &fakeLookup{}
|
||
|
|
p := &ServerPool{resolvers: []Resolver{a, b}}
|
||
|
|
if _, err := p.LookupIP(context.Background(), "missing.example.com"); err == nil {
|
||
|
|
t.Fatal("expected error")
|
||
|
|
}
|
||
|
|
if b.calls.Load() != 0 {
|
||
|
|
t.Fatalf("NXDOMAIN must not fail over; b.calls=%d", b.calls.Load())
|
||
|
|
}
|
||
|
|
}
|