use reflect.DeepEqual for error check in tests

grpc · Oct 18, 2022 · f68e5d4 · f68e5d4
1 parent 79ccdd8
commit f68e5d4
Showing 1 changed file with 62 additions and 34 deletions.
diff --git a/balancer/weightedtarget/weightedtarget_test.go b/balancer/weightedtarget/weightedtarget_test.go
@@ -19,10 +19,11 @@
 package weightedtarget
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
-	"strings"
+	"reflect"
 	"testing"
 	"time"
 
@@ -40,6 +41,10 @@ import (
 	"google.golang.org/grpc/serviceconfig"
 )
 
+const (
+	defaultTestTimeout = 5 * time.Second
+)
+
 type s struct {
 	grpctest.Tester
 }
@@ -571,10 +576,10 @@ func (s) TestWeightedTarget_TwoSubBalancers_MoreBackends(t *testing.T) {
 	}
 
 	// 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}
@@ -594,13 +599,21 @@ 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, func(p balancer.Picker) error {
+		for i := 0; i < 5; i++ {
+			if _, err := p.Pick(balancer.PickInfo{}); !reflect.DeepEqual(err, fmt.Errorf("last connection error: subConn connection error")) {
+				return fmt.Errorf("want pick error %q, got error %q", wantSubConnErr, err)
+			}
 		}
+		return nil
+	}); err != nil {
+		t.Fatal(err.Error())
 	}
 }
 
@@ -802,10 +815,10 @@ func (s) TestWeightedTarget_ThreeSubBalancers_RemoveBalancer(t *testing.T) {
 	}
 
 	// 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
 
@@ -833,16 +846,24 @@ 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, func(p balancer.Picker) error {
+		for i := 0; i < 5; i++ {
+			if _, err := p.Pick(balancer.PickInfo{}); !reflect.DeepEqual(err, fmt.Errorf("last connection error: subConn connection error")) {
+				return fmt.Errorf("want pick error %q, got error %q", wantSubConnErr, err)
+			}
 		}
+		return nil
+	}); err != nil {
+		t.Fatal(err.Error())
 	}
 }
 
@@ -1077,38 +1098,45 @@ 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)
-		}
-	}
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
 
-	// Set one subconn to Connecting, it shouldn't change the overall state.
-	wtb.UpdateSubConnState(sc1, balancer.SubConnState{ConnectivityState: connectivity.Connecting})
-	select {
-	case <-time.After(100 * time.Millisecond):
-	case <-cc.NewPickerCh:
-		t.Fatal("received new picker from the LB policy when expecting none")
-	}
+	if err := cc.WaitForPicker(ctx, func(p balancer.Picker) error {
+		for i := 0; i < 5; i++ {
+			if _, err := p.Pick(balancer.PickInfo{}); !reflect.DeepEqual(err, fmt.Errorf("last connection error: subConn connection error")) {
+				return fmt.Errorf("want pick error %q, got error %q", wantSubConnErr, err)
+			}
+		}
+		// Set one subconn to Connecting, it shouldn't change the overall state.
+		wtb.UpdateSubConnState(sc1, balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+		select {
+		case <-time.After(100 * time.Millisecond):
+		case <-cc.NewPickerCh:
+			t.Fatal("received new picker from the LB policy when expecting none")
+		}
 
-	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)
+		// TODO: look here
+		for i := 0; i < 5; i++ {
+			if _, err := p.Pick(balancer.PickInfo{}); !reflect.DeepEqual(err, fmt.Errorf("last connection error: subConn connection error")) {
+				t.Fatalf("want pick error %q, got err %q", wantSubConnErr, err)
+			}
 		}
+		return nil
+	}); err != nil {
+		t.Fatal(err.Error())
 	}
+
 }
 
 // Verify that a SubConn is created with the expected address and hierarchy