127 lines
3.9 KiB
Go
127 lines
3.9 KiB
Go
package noncebalancer
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
"google.golang.org/grpc/balancer"
|
|
"google.golang.org/grpc/balancer/base"
|
|
"google.golang.org/grpc/resolver"
|
|
|
|
"github.com/letsencrypt/boulder/nonce"
|
|
"github.com/letsencrypt/boulder/test"
|
|
)
|
|
|
|
func TestPickerPicksCorrectBackend(t *testing.T) {
|
|
_, p, subConns := setupTest(false)
|
|
prefix := nonce.DerivePrefix(subConns[0].addrs[0].Addr, []byte("Kala namak"))
|
|
|
|
testCtx := context.WithValue(context.Background(), nonce.PrefixCtxKey{}, "HNmOnt8w")
|
|
testCtx = context.WithValue(testCtx, nonce.HMACKeyCtxKey{}, []byte(prefix))
|
|
info := balancer.PickInfo{Ctx: testCtx}
|
|
|
|
gotPick, err := p.Pick(info)
|
|
test.AssertNotError(t, err, "Pick failed")
|
|
test.AssertDeepEquals(t, subConns[0], gotPick.SubConn)
|
|
}
|
|
|
|
func TestPickerMissingPrefixInCtx(t *testing.T) {
|
|
_, p, subConns := setupTest(false)
|
|
prefix := nonce.DerivePrefix(subConns[0].addrs[0].Addr, []byte("Kala namak"))
|
|
|
|
testCtx := context.WithValue(context.Background(), nonce.HMACKeyCtxKey{}, []byte(prefix))
|
|
info := balancer.PickInfo{Ctx: testCtx}
|
|
|
|
gotPick, err := p.Pick(info)
|
|
test.AssertErrorIs(t, err, errMissingPrefixCtxKey)
|
|
test.AssertNil(t, gotPick.SubConn, "subConn should be nil")
|
|
}
|
|
|
|
func TestPickerInvalidPrefixInCtx(t *testing.T) {
|
|
_, p, _ := setupTest(false)
|
|
|
|
testCtx := context.WithValue(context.Background(), nonce.PrefixCtxKey{}, 9)
|
|
testCtx = context.WithValue(testCtx, nonce.HMACKeyCtxKey{}, []byte("foobar"))
|
|
info := balancer.PickInfo{Ctx: testCtx}
|
|
|
|
gotPick, err := p.Pick(info)
|
|
test.AssertErrorIs(t, err, errInvalidPrefixCtxKeyType)
|
|
test.AssertNil(t, gotPick.SubConn, "subConn should be nil")
|
|
}
|
|
|
|
func TestPickerMissingHMACKeyInCtx(t *testing.T) {
|
|
_, p, _ := setupTest(false)
|
|
|
|
testCtx := context.WithValue(context.Background(), nonce.PrefixCtxKey{}, "HNmOnt8w")
|
|
info := balancer.PickInfo{Ctx: testCtx}
|
|
|
|
gotPick, err := p.Pick(info)
|
|
test.AssertErrorIs(t, err, errMissingHMACKeyCtxKey)
|
|
test.AssertNil(t, gotPick.SubConn, "subConn should be nil")
|
|
}
|
|
|
|
func TestPickerInvalidHMACKeyInCtx(t *testing.T) {
|
|
_, p, _ := setupTest(false)
|
|
|
|
testCtx := context.WithValue(context.Background(), nonce.PrefixCtxKey{}, "HNmOnt8w")
|
|
testCtx = context.WithValue(testCtx, nonce.HMACKeyCtxKey{}, 9)
|
|
info := balancer.PickInfo{Ctx: testCtx}
|
|
|
|
gotPick, err := p.Pick(info)
|
|
test.AssertErrorIs(t, err, errInvalidHMACKeyCtxKeyType)
|
|
test.AssertNil(t, gotPick.SubConn, "subConn should be nil")
|
|
}
|
|
|
|
func TestPickerNoMatchingSubConnAvailable(t *testing.T) {
|
|
_, p, subConns := setupTest(false)
|
|
prefix := nonce.DerivePrefix(subConns[0].addrs[0].Addr, []byte("Kala namak"))
|
|
|
|
testCtx := context.WithValue(context.Background(), nonce.PrefixCtxKey{}, "rUsTrUin")
|
|
testCtx = context.WithValue(testCtx, nonce.HMACKeyCtxKey{}, []byte(prefix))
|
|
info := balancer.PickInfo{Ctx: testCtx}
|
|
|
|
gotPick, err := p.Pick(info)
|
|
test.AssertErrorIs(t, err, ErrNoBackendsMatchPrefix.Err())
|
|
test.AssertNil(t, gotPick.SubConn, "subConn should be nil")
|
|
}
|
|
|
|
func TestPickerNoSubConnsAvailable(t *testing.T) {
|
|
b, p, _ := setupTest(true)
|
|
b.Build(base.PickerBuildInfo{})
|
|
info := balancer.PickInfo{Ctx: context.Background()}
|
|
|
|
gotPick, err := p.Pick(info)
|
|
test.AssertErrorIs(t, err, balancer.ErrNoSubConnAvailable)
|
|
test.AssertNil(t, gotPick.SubConn, "subConn should be nil")
|
|
}
|
|
|
|
func setupTest(noSubConns bool) (*Balancer, balancer.Picker, []*subConn) {
|
|
var subConns []*subConn
|
|
bi := base.PickerBuildInfo{
|
|
ReadySCs: make(map[balancer.SubConn]base.SubConnInfo),
|
|
}
|
|
|
|
sc := &subConn{}
|
|
addr := resolver.Address{Addr: "10.77.77.77:8080"}
|
|
sc.UpdateAddresses([]resolver.Address{addr})
|
|
|
|
if !noSubConns {
|
|
bi.ReadySCs[sc] = base.SubConnInfo{Address: addr}
|
|
subConns = append(subConns, sc)
|
|
}
|
|
|
|
b := &Balancer{}
|
|
p := b.Build(bi)
|
|
return b, p, subConns
|
|
}
|
|
|
|
// subConn is a test mock which implements the balancer.SubConn interface.
|
|
type subConn struct {
|
|
balancer.SubConn
|
|
addrs []resolver.Address
|
|
}
|
|
|
|
func (s *subConn) UpdateAddresses(addrs []resolver.Address) {
|
|
s.addrs = addrs
|
|
}
|