Review comments #2.

easwars · Aug 27, 2020 · deaaaf0 · deaaaf0
1 parent 6fc6be3
commit deaaaf0
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 7 deletions.
diff --git a/xds/internal/balancer/edsbalancer/xds_client_wrapper_test.go b/xds/internal/balancer/edsbalancer/xds_client_wrapper_test.go
@@ -167,7 +167,10 @@ func (s) TestClientWrapperHandleUpdateError(t *testing.T) {
 
 	xdsC := fakeclient.NewClient()
 	cw.handleUpdate(&EDSConfig{EDSServiceName: testEDSClusterName}, attributes.New(xdsinternal.XDSClientID, xdsC))
-	gotCluster, err := xdsC.WaitForWatchEDS(context.Background())
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	gotCluster, err := xdsC.WaitForWatchEDS(ctx)
 	if err != nil {
 		t.Fatalf("xdsClient.WatchEndpoints failed with error: %v", err)
 	}
@@ -181,7 +184,7 @@ func (s) TestClientWrapperHandleUpdateError(t *testing.T) {
 	//
 	// TODO: check for loseContact() when errors indicating "lose contact" are
 	// handled correctly.
-	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	ctx, cancel = context.WithTimeout(context.Background(), defaultTestTimeout)
 	defer cancel()
 	gotUpdate, err := edsRespChan.Receive(ctx)
 	if err != nil {

diff --git a/xds/internal/balancer/edsbalancer/xds_lrs_test.go b/xds/internal/balancer/edsbalancer/xds_lrs_test.go
@@ -19,7 +19,7 @@
 package edsbalancer
 
 import (
-	context2 "context"
+	"context"
 	"testing"
 
 	"google.golang.org/grpc/attributes"
@@ -42,20 +42,24 @@ func (s) TestXDSLoadReporting(t *testing.T) {
 	defer edsB.Close()
 
 	xdsC := fakeclient.NewClient()
-	edsB.UpdateClientConnState(balancer.ClientConnState{
+	if err := edsB.UpdateClientConnState(balancer.ClientConnState{
 		ResolverState:  resolver.State{Attributes: attributes.New(xdsinternal.XDSClientID, xdsC)},
 		BalancerConfig: &EDSConfig{LrsLoadReportingServerName: new(string)},
-	})
+	}); err != nil {
+		t.Fatal(err)
+	}
 
-	gotCluster, err := xdsC.WaitForWatchEDS(context2.Background())
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	gotCluster, err := xdsC.WaitForWatchEDS(ctx)
 	if err != nil {
 		t.Fatalf("xdsClient.WatchEndpoints failed with error: %v", err)
 	}
 	if gotCluster != testEDSClusterName {
 		t.Fatalf("xdsClient.WatchEndpoints() called with cluster: %v, want %v", gotCluster, testEDSClusterName)
 	}
 
-	got, err := xdsC.WaitForReportLoad(context2.Background())
+	got, err := xdsC.WaitForReportLoad(ctx)
 	if err != nil {
 		t.Fatalf("xdsClient.ReportLoad failed with error: %v", err)
 	}