fix flaky test

Author: Pranjali-2501

internal/xds/balancer/cdsbalancer/aggregate_cluster_test.go

@@ -18,6 +18,7 @@ package cdsbalancer
 
 import (
 	"context"
+	"fmt"
 	"strings"
 	"testing"
 	"time"
@@ -31,6 +32,7 @@ import (
 	"google.golang.org/grpc/internal/testutils/xds/e2e"
 	"google.golang.org/grpc/internal/xds/balancer/priority"
 	"google.golang.org/grpc/internal/xds/xdsclient/xdsresource/version"
+	"google.golang.org/grpc/resolver"
 	"google.golang.org/grpc/serviceconfig"
 	"google.golang.org/grpc/status"
 
@@ -167,6 +169,7 @@ func (s) TestAggregateClusterSuccess_LeafNode(t *testing.T) {
 // LogicalDNS and verifies that the load balancing configuration pushed to the
 // priority LB policy contains the expected config.
 func (s) TestAggregateClusterSuccess_ThenUpdateChildClusters(t *testing.T) {
+	dnsTargetCh, dnsR := setupDNS(t)
 	lbCfgCh, _, _, _ := registerWrappedPriorityPolicy(t)
 	mgmtServer, nodeID, _ := setupWithManagementServer(t, nil, nil)
 
@@ -205,6 +208,20 @@ func (s) TestAggregateClusterSuccess_ThenUpdateChildClusters(t *testing.T) {
 		t.Fatal(err)
 	}
 
+	// Ensure that the DNS resolver is started for the expected target.
+	select {
+	case <-ctx.Done():
+		t.Fatal("Timeout waiting for DNS watch")
+	case target := <-dnsTargetCh:
+		addr := fmt.Sprintf("%s:%d", dnsHostName, dnsPort)
+		if target.Endpoint() != addr {
+			t.Fatalf("DNS resolution started for target %q, want %q", target.Endpoint(), addr)
+		}
+		dnsR.UpdateState(resolver.State{
+			Endpoints: []resolver.Endpoint{{Addresses: []resolver.Address{{Addr: addr}}}},
+		})
+	}
+
 	wantChildCfg := &priority.LBConfig{
 		Children: map[string]*priority.Child{
 			"priority-0-0": {
@@ -236,6 +253,18 @@ func (s) TestAggregateClusterSuccess_ThenUpdateChildClusters(t *testing.T) {
 	if err := mgmtServer.Update(ctx, resources); err != nil {
 		t.Fatal(err)
 	}
+
+	select {
+	case <-ctx.Done():
+		t.Fatal("Timeout waiting for DNS watch")
+	case <-dnsTargetCh:
+		addr1 := fmt.Sprintf("%s:%d", dnsHostName, dnsPort)
+		addr2 := fmt.Sprintf("%s:%d", dnsHostNameNew, dnsPort)
+		dnsR.UpdateState(resolver.State{
+			Endpoints: []resolver.Endpoint{{Addresses: []resolver.Address{{Addr: addr1}, {Addr: addr2}}}},
+		})
+	}
+
 	wantChildCfg = &priority.LBConfig{
 		Children: map[string]*priority.Child{
 			"priority-0-0": {
@@ -260,6 +289,7 @@ func (s) TestAggregateClusterSuccess_ThenUpdateChildClusters(t *testing.T) {
 // configuration pushed to the priority LB policy contains a single discovery
 // mechanism.
 func (s) TestAggregateClusterSuccess_ThenChangeRootToEDS(t *testing.T) {
+	dnsTargetCh, dnsR := setupDNS(t)
 	lbCfgCh, _, _, _ := registerWrappedPriorityPolicy(t)
 	mgmtServer, nodeID, _ := setupWithManagementServer(t, nil, nil)
 
@@ -282,6 +312,20 @@ func (s) TestAggregateClusterSuccess_ThenChangeRootToEDS(t *testing.T) {
 		t.Fatal(err)
 	}
 
+	// Ensure that the DNS resolver is started for the expected target.
+	select {
+	case <-ctx.Done():
+		t.Fatal("Timeout waiting for DNS watch")
+	case target := <-dnsTargetCh:
+		addr := fmt.Sprintf("%s:%d", dnsHostName, dnsPort)
+		if target.Endpoint() != addr {
+			t.Fatalf("DNS resolution started for target %q, want %q", target.Endpoint(), addr)
+		}
+		dnsR.UpdateState(resolver.State{
+			Endpoints: []resolver.Endpoint{{Addresses: []resolver.Address{{Addr: addr}}}},
+		})
+	}
+
 	wantChildCfg := &priority.LBConfig{
 		Children: map[string]*priority.Child{
 			"priority-0-0": {
@@ -330,6 +374,7 @@ func (s) TestAggregateClusterSuccess_ThenChangeRootToEDS(t *testing.T) {
 // cluster. In each of these cases, the test verifies that the load balancing
 // configuration pushed to the priority LB policy contains the expected config.
 func (s) TestAggregatedClusterSuccess_SwitchBetweenLeafAndAggregate(t *testing.T) {
+	dnsTargetCh, dnsR := setupDNS(t)
 	lbCfgCh, _, _, _ := registerWrappedPriorityPolicy(t)
 	mgmtServer, nodeID, _ := setupWithManagementServer(t, nil, nil)
 
@@ -375,6 +420,21 @@ func (s) TestAggregatedClusterSuccess_SwitchBetweenLeafAndAggregate(t *testing.T
 	if err := mgmtServer.Update(ctx, resources); err != nil {
 		t.Fatal(err)
 	}
+
+	// Ensure that the DNS resolver is started for the expected target.
+	select {
+	case <-ctx.Done():
+		t.Fatal("Timeout waiting for DNS watch")
+	case target := <-dnsTargetCh:
+		addr := fmt.Sprintf("%s:%d", dnsHostName, dnsPort)
+		if target.Endpoint() != addr {
+			t.Fatalf("DNS resolution started for target %q, want %q", target.Endpoint(), addr)
+		}
+		dnsR.UpdateState(resolver.State{
+			Endpoints: []resolver.Endpoint{{Addresses: []resolver.Address{{Addr: addr}}}},
+		})
+	}
+
 	wantChildCfg = &priority.LBConfig{
 		Children: map[string]*priority.Child{
 			"priority-0-0": {

internal/xds/balancer/cdsbalancer/cdsbalancer_test.go

@@ -81,7 +81,7 @@ const (
 	host                    = "localhost"
 	port                    = 8080
 	dnsPort                 = uint32(8080)
-	defaultTestTimeout      = 5 * time.Second
+	defaultTestTimeout      = 10 * time.Second
 	defaultTestShortTimeout = 10 * time.Millisecond // For events expected to *not* happen.
 )
 
@@ -164,6 +164,27 @@ func registerWrappedPriorityPolicy(t *testing.T) (chan serviceconfig.LoadBalanci
 	return lbCfgCh, resolverErrCh, exitIdleCh, closeCh
 }
 
+// setupDNS unregisters the DNS resolver and registers a manual resolver for the
+// same scheme. This allows the test to mock the DNS resolution by supplying the
+// addresses of the test backends.
+//
+// Returns the following:
+//   - a channel onto which the DNS target being resolved is written to by the
+//     mock DNS resolver
+//   - a manual resolver which is used to mock the actual DNS resolution
+func setupDNS(t *testing.T) (chan resolver.Target, *manual.Resolver) {
+	targetCh := make(chan resolver.Target, 1)
+
+	mr := manual.NewBuilderWithScheme("dns")
+	mr.BuildCallback = func(target resolver.Target, _ resolver.ClientConn, _ resolver.BuildOptions) { targetCh <- target }
+
+	dnsResolverBuilder := resolver.Get("dns")
+	resolver.Register(mr)
+
+	t.Cleanup(func() { resolver.Register(dnsResolverBuilder) })
+	return targetCh, mr
+}
+
 // Performs the following setup required for tests:
 //   - Spins up an xDS management server and the provided onStreamRequest
 //     function is set to be called for every incoming request on the ADS stream.