Skip to content

Commit

Permalink
test: refactor test to use testutils helpers
Browse files Browse the repository at this point in the history
  • Loading branch information
@arvindbr8
arvindbr8 committed Oct 18, 2022
1 parent dbb8e2b commit 250b421
Show file tree
Hide file tree
Showing 5 changed files with 135 additions and 177 deletions.
91 changes: 51 additions & 40 deletions balancer/weightedtarget/weightedtarget_test.go
View file
Expand Up @@ -19,6 +19,7 @@
package weightedtarget

import (
"context"
"encoding/json"
"errors"
"fmt"
Expand All @@ -40,6 +41,10 @@ import (
"google.golang.org/grpc/serviceconfig"
)

const (
defaultTestTimeout = 5 * time.Second
)

type s struct {
grpctest.Tester
}
Expand All @@ -58,6 +63,17 @@ func newTestConfigBalancerBuilder() *testConfigBalancerBuilder {
}
}

func makeRPCsAndAssertContainsError(rpcCount int, want error) func(balancer.Picker) error {
return func(p balancer.Picker) error {
for i := 0; i < rpcCount; i++ {
if _, err := p.Pick(balancer.PickInfo{}); !strings.Contains(err.Error(), want.Error()) {
return fmt.Errorf("want pick error to contain %q, got %q", want, err)
}
}
return nil
}
}

func (t *testConfigBalancerBuilder) Build(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
rr := t.Builder.Build(cc, opts)
return &testConfigBalancer{
Expand Down Expand Up @@ -289,13 +305,6 @@ func (s) TestWeightedTarget(t *testing.T) {
}
}

func subConnFromPicker(p balancer.Picker) func() balancer.SubConn {
return func() balancer.SubConn {
scst, _ := p.Pick(balancer.PickInfo{})
return scst.SubConn
}
}

// TestWeightedTarget_OneSubBalancer_AddRemoveBackend tests the case where we
// have a weighted target balancer will one sub-balancer, and we add and remove
// backends from the subBalancer.
Expand Down Expand Up @@ -366,7 +375,7 @@ func (s) TestWeightedTarget_OneSubBalancer_AddRemoveBackend(t *testing.T) {

// Test round robin pick.
want := []balancer.SubConn{sc1, sc2}
if err := testutils.IsRoundRobin(want, subConnFromPicker(p)); err != nil {
if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p)); err != nil {
t.Fatalf("want %v, got %v", want, err)
}

Expand Down Expand Up @@ -455,7 +464,7 @@ func (s) TestWeightedTarget_TwoSubBalancers_OneBackend(t *testing.T) {

// Test roundrobin on the last picker.
want := []balancer.SubConn{sc1, sc2}
if err := testutils.IsRoundRobin(want, subConnFromPicker(p)); err != nil {
if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p)); err != nil {
t.Fatalf("want %v, got %v", want, err)
}
}
Expand Down Expand Up @@ -536,15 +545,15 @@ func (s) TestWeightedTarget_TwoSubBalancers_MoreBackends(t *testing.T) {
// Test roundrobin on the last picker. RPCs should be sent equally to all
// backends.
want := []balancer.SubConn{sc1, sc2, sc3, sc4}
if err := testutils.IsRoundRobin(want, subConnFromPicker(p)); err != nil {
if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p)); err != nil {
t.Fatalf("want %v, got %v", want, err)
}

// Turn sc2's connection down, should be RR between balancers.
wtb.UpdateSubConnState(sc2, balancer.SubConnState{ConnectivityState: connectivity.TransientFailure})
p = <-cc.NewPickerCh
want = []balancer.SubConn{sc1, sc1, sc3, sc4}
if err := testutils.IsRoundRobin(want, subConnFromPicker(p)); err != nil {
if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p)); err != nil {
t.Fatalf("want %v, got %v", want, err)
}

Expand All @@ -566,19 +575,19 @@ func (s) TestWeightedTarget_TwoSubBalancers_MoreBackends(t *testing.T) {
wtb.UpdateSubConnState(scRemoved, balancer.SubConnState{ConnectivityState: connectivity.Shutdown})
p = <-cc.NewPickerCh
want = []balancer.SubConn{sc1, sc4}
if err := testutils.IsRoundRobin(want, subConnFromPicker(p)); err != nil {
if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p)); err != nil {
t.Fatalf("want %v, got %v", want, err)
}

// Turn sc1's connection down.
scConnErr := errors.New("subConn connection error")
wantSubConnErr := errors.New("subConn connection error")
wtb.UpdateSubConnState(sc1, balancer.SubConnState{
ConnectivityState: connectivity.TransientFailure,
ConnectionError: scConnErr,
ConnectionError: wantSubConnErr,
})
p = <-cc.NewPickerCh
want = []balancer.SubConn{sc4}
if err := testutils.IsRoundRobin(want, subConnFromPicker(p)); err != nil {
if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p)); err != nil {
t.Fatalf("want %v, got %v", want, err)
}

Expand All @@ -594,13 +603,14 @@ func (s) TestWeightedTarget_TwoSubBalancers_MoreBackends(t *testing.T) {
// Turn all connections down.
wtb.UpdateSubConnState(sc4, balancer.SubConnState{
ConnectivityState: connectivity.TransientFailure,
ConnectionError: scConnErr,
ConnectionError: wantSubConnErr,
})
p = <-cc.NewPickerCh
for i := 0; i < 5; i++ {
if _, err := p.Pick(balancer.PickInfo{}); err == nil || !strings.Contains(err.Error(), scConnErr.Error()) {
t.Fatalf("want pick error %q, got error %q", scConnErr, err)
}

ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()

if err := cc.WaitForPicker(ctx, makeRPCsAndAssertContainsError(5, wantSubConnErr)); err != nil {
t.Fatal(err.Error())
}
}

Expand Down Expand Up @@ -680,7 +690,7 @@ func (s) TestWeightedTarget_TwoSubBalancers_DifferentWeight_MoreBackends(t *test
// Test roundrobin on the last picker. Twice the number of RPCs should be
// sent to cluster_1 when compared to cluster_2.
want := []balancer.SubConn{sc1, sc1, sc2, sc2, sc3, sc4}
if err := testutils.IsRoundRobin(want, subConnFromPicker(p)); err != nil {
if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p)); err != nil {
t.Fatalf("want %v, got %v", want, err)
}
}
Expand Down Expand Up @@ -757,7 +767,7 @@ func (s) TestWeightedTarget_ThreeSubBalancers_RemoveBalancer(t *testing.T) {
p := <-cc.NewPickerCh

want := []balancer.SubConn{sc1, sc2, sc3}
if err := testutils.IsRoundRobin(want, subConnFromPicker(p)); err != nil {
if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p)); err != nil {
t.Fatalf("want %v, got %v", want, err)
}

Expand Down Expand Up @@ -797,15 +807,15 @@ func (s) TestWeightedTarget_ThreeSubBalancers_RemoveBalancer(t *testing.T) {
t.Fatalf("RemoveSubConn, want %v, got %v", sc2, scRemoved)
}
want = []balancer.SubConn{sc1, sc3}
if err := testutils.IsRoundRobin(want, subConnFromPicker(p)); err != nil {
if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p)); err != nil {
t.Fatalf("want %v, got %v", want, err)
}

// Move balancer 3 into transient failure.
scConnErr := errors.New("subConn connection error")
wantSubConnErr := errors.New("subConn connection error")
wtb.UpdateSubConnState(sc3, balancer.SubConnState{
ConnectivityState: connectivity.TransientFailure,
ConnectionError: scConnErr,
ConnectionError: wantSubConnErr,
})
<-cc.NewPickerCh

Expand Down Expand Up @@ -833,16 +843,17 @@ func (s) TestWeightedTarget_ThreeSubBalancers_RemoveBalancer(t *testing.T) {

// Removing a subBalancer causes the weighted target LB policy to push a new
// picker which ensures that the removed subBalancer is not picked for RPCs.
p = <-cc.NewPickerCh

scRemoved = <-cc.RemoveSubConnCh
if !cmp.Equal(scRemoved, sc1, cmp.AllowUnexported(testutils.TestSubConn{})) {
t.Fatalf("RemoveSubConn, want %v, got %v", sc1, scRemoved)
}
for i := 0; i < 5; i++ {
if _, err := p.Pick(balancer.PickInfo{}); err == nil || !strings.Contains(err.Error(), scConnErr.Error()) {
t.Fatalf("want pick error %q, got error %q", scConnErr, err)
}

ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
defer cancel()

if err := cc.WaitForPicker(ctx, makeRPCsAndAssertContainsError(5, wantSubConnErr)); err != nil {
t.Fatal(err.Error())
}
}

Expand Down Expand Up @@ -922,7 +933,7 @@ func (s) TestWeightedTarget_TwoSubBalancers_ChangeWeight_MoreBackends(t *testing
// Test roundrobin on the last picker. Twice the number of RPCs should be
// sent to cluster_1 when compared to cluster_2.
want := []balancer.SubConn{sc1, sc1, sc2, sc2, sc3, sc4}
if err := testutils.IsRoundRobin(want, subConnFromPicker(p)); err != nil {
if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p)); err != nil {
t.Fatalf("want %v, got %v", want, err)
}

Expand Down Expand Up @@ -958,7 +969,7 @@ func (s) TestWeightedTarget_TwoSubBalancers_ChangeWeight_MoreBackends(t *testing
// Weight change causes a new picker to be pushed to the channel.
p = <-cc.NewPickerCh
want = []balancer.SubConn{sc1, sc1, sc1, sc2, sc2, sc2, sc3, sc4}
if err := testutils.IsRoundRobin(want, subConnFromPicker(p)); err != nil {
if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p)); err != nil {
t.Fatalf("want %v, got %v", want, err)
}
}
Expand Down Expand Up @@ -1077,21 +1088,21 @@ func (s) TestBalancerGroup_SubBalancerTurnsConnectingFromTransientFailure(t *tes

// Set both subconn to TransientFailure, this will put both sub-balancers in
// transient failure.
scConnErr := errors.New("subConn connection error")
wantSubConnErr := errors.New("subConn connection error")
wtb.UpdateSubConnState(sc1, balancer.SubConnState{
ConnectivityState: connectivity.TransientFailure,
ConnectionError: scConnErr,
ConnectionError: wantSubConnErr,
})
<-cc.NewPickerCh
wtb.UpdateSubConnState(sc2, balancer.SubConnState{
ConnectivityState: connectivity.TransientFailure,
ConnectionError: scConnErr,
ConnectionError: wantSubConnErr,
})
p := <-cc.NewPickerCh

for i := 0; i < 5; i++ {
if _, err := p.Pick(balancer.PickInfo{}); err == nil || !strings.Contains(err.Error(), scConnErr.Error()) {
t.Fatalf("want pick error %q, got error %q", scConnErr, err)
if _, err := p.Pick(balancer.PickInfo{}); err == nil || !strings.Contains(err.Error(), wantSubConnErr.Error()) {
t.Fatalf("want pick error %q, got error %q", wantSubConnErr, err)
}
}

Expand All @@ -1105,8 +1116,8 @@ func (s) TestBalancerGroup_SubBalancerTurnsConnectingFromTransientFailure(t *tes

for i := 0; i < 5; i++ {
r, err := p.Pick(balancer.PickInfo{})
if err == nil || !strings.Contains(err.Error(), scConnErr.Error()) {
t.Fatalf("want pick error %q, got result %v, err %q", scConnErr, r, err)
if err == nil || !strings.Contains(err.Error(), wantSubConnErr.Error()) {
t.Fatalf("want pick error %q, got result %v, err %q", wantSubConnErr, r, err)
}
}
}
Expand Down
21 changes: 7 additions & 14 deletions internal/balancergroup/balancergroup_test.go
View file
Expand Up @@ -67,13 +67,6 @@ func Test(t *testing.T) {
grpctest.RunSubTests(t, s{})
}

func subConnFromPicker(p balancer.Picker) func() balancer.SubConn {
return func() balancer.SubConn {
scst, _ := p.Pick(balancer.PickInfo{})
return scst.SubConn
}
}

// Create a new balancer group, add balancer and backends, but not start.
// - b1, weight 2, backends [0,1]
// - b2, weight 1, backends [2,3]
Expand Down Expand Up @@ -116,7 +109,7 @@ func (s) TestBalancerGroup_start_close(t *testing.T) {
m1[testBackendAddrs[1]], m1[testBackendAddrs[1]],
m1[testBackendAddrs[2]], m1[testBackendAddrs[3]],
}
if err := testutils.IsRoundRobin(want, subConnFromPicker(p1)); err != nil {
if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p1)); err != nil {
t.Fatalf("want %v, got %v", want, err)
}

Expand Down Expand Up @@ -158,7 +151,7 @@ func (s) TestBalancerGroup_start_close(t *testing.T) {
m2[testBackendAddrs[3]], m2[testBackendAddrs[3]], m2[testBackendAddrs[3]],
m2[testBackendAddrs[1]], m2[testBackendAddrs[2]],
}
if err := testutils.IsRoundRobin(want, subConnFromPicker(p2)); err != nil {
if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p2)); err != nil {
t.Fatalf("want %v, got %v", want, err)
}
}
Expand Down Expand Up @@ -241,7 +234,7 @@ func initBalancerGroupForCachingTest(t *testing.T) (*weightedaggregator.Aggregat
m1[testBackendAddrs[1]], m1[testBackendAddrs[1]],
m1[testBackendAddrs[2]], m1[testBackendAddrs[3]],
}
if err := testutils.IsRoundRobin(want, subConnFromPicker(p1)); err != nil {
if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p1)); err != nil {
t.Fatalf("want %v, got %v", want, err)
}

Expand All @@ -263,7 +256,7 @@ func initBalancerGroupForCachingTest(t *testing.T) (*weightedaggregator.Aggregat
want = []balancer.SubConn{
m1[testBackendAddrs[0]], m1[testBackendAddrs[1]],
}
if err := testutils.IsRoundRobin(want, subConnFromPicker(p2)); err != nil {
if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p2)); err != nil {
t.Fatalf("want %v, got %v", want, err)
}

Expand Down Expand Up @@ -304,7 +297,7 @@ func (s) TestBalancerGroup_locality_caching(t *testing.T) {
// addr2 is down, b2 only has addr3 in READY state.
addrToSC[testBackendAddrs[3]], addrToSC[testBackendAddrs[3]],
}
if err := testutils.IsRoundRobin(want, subConnFromPicker(p3)); err != nil {
if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p3)); err != nil {
t.Fatalf("want %v, got %v", want, err)
}

Expand Down Expand Up @@ -452,7 +445,7 @@ func (s) TestBalancerGroup_locality_caching_readd_with_different_builder(t *test
addrToSC[testBackendAddrs[1]], addrToSC[testBackendAddrs[1]],
addrToSC[testBackendAddrs[4]], addrToSC[testBackendAddrs[5]],
}
if err := testutils.IsRoundRobin(want, subConnFromPicker(p3)); err != nil {
if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p3)); err != nil {
t.Fatalf("want %v, got %v", want, err)
}
}
Expand Down Expand Up @@ -583,7 +576,7 @@ func (s) TestBalancerGracefulSwitch(t *testing.T) {
want := []balancer.SubConn{
m1[testBackendAddrs[0]], m1[testBackendAddrs[1]],
}
if err := testutils.IsRoundRobin(want, subConnFromPicker(p1)); err != nil {
if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p1)); err != nil {
t.Fatalf("want %v, got %v", want, err)
}

Expand Down
7 changes: 7 additions & 0 deletions internal/testutils/balancer.go
View file
Expand Up @@ -339,6 +339,13 @@ func IsRoundRobin(want []balancer.SubConn, f func() balancer.SubConn) error {
return nil
}

func SubConnFromPicker(p balancer.Picker) func() balancer.SubConn {
return func() balancer.SubConn {
scst, _ := p.Pick(balancer.PickInfo{})
return scst.SubConn
}
}

// ErrTestConstPicker is error returned by test const picker.
var ErrTestConstPicker = fmt.Errorf("const picker error")

Expand Down

0 comments on commit 250b421

Please sign in to comment.