105 lines
3.1 KiB
Go
105 lines
3.1 KiB
Go
package loadbalance
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"github.com/miekg/coredns/middleware"
|
|
"github.com/miekg/coredns/middleware/test"
|
|
|
|
"github.com/miekg/dns"
|
|
"golang.org/x/net/context"
|
|
)
|
|
|
|
func TestLoadBalance(t *testing.T) {
|
|
rm := RoundRobin{Next: handler()}
|
|
|
|
// the first X records must be cnames after this test
|
|
tests := []struct {
|
|
answer []dns.RR
|
|
extra []dns.RR
|
|
cnameAnswer int
|
|
cnameExtra int
|
|
}{
|
|
{
|
|
answer: []dns.RR{
|
|
newCNAME("cname1.region2.skydns.test. 300 IN CNAME cname2.region2.skydns.test."),
|
|
newCNAME("cname2.region2.skydns.test. 300 IN CNAME cname3.region2.skydns.test."),
|
|
newCNAME("cname5.region2.skydns.test. 300 IN CNAME cname6.region2.skydns.test."),
|
|
newCNAME("cname6.region2.skydns.test. 300 IN CNAME endpoint.region2.skydns.test."),
|
|
newA("endpoint.region2.skydns.test. 300 IN A 10.240.0.1"),
|
|
},
|
|
cnameAnswer: 4,
|
|
},
|
|
{
|
|
answer: []dns.RR{
|
|
newA("endpoint.region2.skydns.test. 300 IN A 10.240.0.1"),
|
|
newCNAME("cname.region2.skydns.test. 300 IN CNAME endpoint.region2.skydns.test."),
|
|
},
|
|
cnameAnswer: 1,
|
|
},
|
|
{
|
|
answer: []dns.RR{
|
|
newA("endpoint.region2.skydns.test. 300 IN A 10.240.0.1"),
|
|
newA("endpoint.region2.skydns.test. 300 IN A 10.240.0.2"),
|
|
newCNAME("cname2.region2.skydns.test. 300 IN CNAME cname3.region2.skydns.test."),
|
|
newA("endpoint.region2.skydns.test. 300 IN A 10.240.0.3"),
|
|
},
|
|
extra: []dns.RR{
|
|
newA("endpoint.region2.skydns.test. 300 IN A 10.240.0.1"),
|
|
newAAAA("endpoint.region2.skydns.test. 300 IN AAAA ::1"),
|
|
newCNAME("cname2.region2.skydns.test. 300 IN CNAME cname3.region2.skydns.test."),
|
|
newA("endpoint.region2.skydns.test. 300 IN A 10.240.0.3"),
|
|
newAAAA("endpoint.region2.skydns.test. 300 IN AAAA ::2"),
|
|
},
|
|
cnameAnswer: 1,
|
|
cnameExtra: 1,
|
|
},
|
|
}
|
|
|
|
rec := middleware.NewResponseRecorder(&test.ResponseWriter{})
|
|
|
|
for i, test := range tests {
|
|
req := new(dns.Msg)
|
|
req.SetQuestion("region2.skydns.test.", dns.TypeSRV)
|
|
req.Answer = test.answer
|
|
req.Extra = test.extra
|
|
|
|
_, err := rm.ServeDNS(context.TODO(), rec, req)
|
|
if err != nil {
|
|
t.Errorf("Test %d: Expected no error, but got %s", i, err)
|
|
continue
|
|
|
|
}
|
|
cname := 0
|
|
for _, r := range rec.Msg().Answer {
|
|
if r.Header().Rrtype != dns.TypeCNAME {
|
|
break
|
|
}
|
|
cname++
|
|
}
|
|
if cname != test.cnameAnswer {
|
|
t.Errorf("Test %d: Expected %d cnames in Answer, but got %d", i, test.cnameAnswer, cname)
|
|
}
|
|
cname = 0
|
|
for _, r := range rec.Msg().Extra {
|
|
if r.Header().Rrtype != dns.TypeCNAME {
|
|
break
|
|
}
|
|
cname++
|
|
}
|
|
if cname != test.cnameExtra {
|
|
t.Errorf("Test %d: Expected %d cname in Extra, but got %d", i, test.cnameExtra, cname)
|
|
}
|
|
}
|
|
}
|
|
|
|
func handler() middleware.Handler {
|
|
return middleware.HandlerFunc(func(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
|
|
w.WriteMsg(r)
|
|
return dns.RcodeSuccess, nil
|
|
})
|
|
}
|
|
|
|
func newA(rr string) *dns.A { r, _ := dns.NewRR(rr); return r.(*dns.A) }
|
|
func newAAAA(rr string) *dns.AAAA { r, _ := dns.NewRR(rr); return r.(*dns.AAAA) }
|
|
func newCNAME(rr string) *dns.CNAME { r, _ := dns.NewRR(rr); return r.(*dns.CNAME) }
|