Repository: grpc/grpc-go
File: xds/internal/balancer/cdsbalancer/aggregate_cluster_test.go (Go)

Code before:
		t.Fatalf("Child policy received configuration when not expected to: %s", pretty.ToJSON(cfg))
	case <-time.After(defaultTestShortTimeout):
	}

	testutils.AwaitState(ctx, t, cc, connectivity.TransientFailure)

	// Verify that the RPC fails with expected code.
	client := testgrpc.NewTestServiceClient(cc)
	_, err := client.EmptyCall(ctx, &testpb.Empty{})
	if gotCode, wantCode := status.Code(err), codes.Unavailable; gotCode != wantCode {
		t.Fatalf("EmptyCall() failed with code: %v, want %v", gotCode, wantCode)
	}
	const wantErr = "aggregate cluster graph has no leaf clusters"
	if !strings.Contains(err.Error(), wantErr) {
		t.Fatalf("EmptyCall() failed with err: %v, want %s", err, wantErr)
	}
}

// Tests the scenario where the aggregate cluster graph contains a cycle and
// also contains a leaf cluster. The case used here is [A -> B, B -> A, C]. As
// there is a leaf cluster in this graph , configuration should be pushed to the
// child policy and RPCs should get routed to that leaf cluster.
func (s) TestAggregatedCluster_CycleWithLeafNode(t *testing.T) {
	lbCfgCh, _, _, _ := registerWrappedClusterResolverPolicy(t)
	mgmtServer, nodeID, cc, _, _, _, _ := setupWithManagementServer(t)
	cc.Connect()
	// Start a test service backend.
	server := stubserver.StartTestService(t, nil)
	t.Cleanup(server.Stop)

	const (
		clusterNameA = clusterName // cluster name in cds LB policy config
		clusterNameB = clusterName + "-B"
		clusterNameC = clusterName + "-C"
	)
	// Configure the management server with an aggregate cluster resource graph
	// that contains a cycle, but also contains a leaf cluster.
	resources := e2e.UpdateOptions{
		NodeID: nodeID,
		Clusters: []*v3clusterpb.Cluster{
			makeAggregateClusterResource(clusterNameA, []string{clusterNameB}),
			makeAggregateClusterResource(clusterNameB, []string{clusterNameA, clusterNameC}),
			e2e.DefaultCluster(clusterNameC, serviceName, e2e.SecurityLevelNone),
		},
		Endpoints:      []*v3endpointpb.ClusterLoadAssignment{e2e.DefaultEndpoint(serviceName, "localhost", []uint32{testutils.ParsePort(t, server.Address)})},
		SkipValidation: true,
	}
	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
	defer cancel()
	if err := mgmtServer.Update(ctx, resources); err != nil {
		t.Fatal(err)

Review comment (refactor):
Do we really need to call `Connect` here? Can this be moved inside the helper function `setupWithManagementServer`?

Code after:
		t.Fatalf("Child policy received configuration when not expected to: %s", pretty.ToJSON(cfg))
	case <-time.After(defaultTestShortTimeout):
	}

	testutils.AwaitState(ctx, t, cc, connectivity.TransientFailure)

	// Verify that the RPC fails with expected code.
	client := testgrpc.NewTestServiceClient(cc)
	_, err := client.EmptyCall(ctx, &testpb.Empty{})
	if gotCode, wantCode := status.Code(err), codes.Unavailable; gotCode != wantCode {
		t.Fatalf("EmptyCall() failed with code: %v, want %v", gotCode, wantCode)
	}
	const wantErr = "aggregate cluster graph has no leaf clusters"
	if !strings.Contains(err.Error(), wantErr) {
		t.Fatalf("EmptyCall() failed with err: %v, want %s", err, wantErr)
	}
}

// Tests the scenario where the aggregate cluster graph contains a cycle and
// also contains a leaf cluster. The case used here is [A -> B, B -> A, C]. As
// there is a leaf cluster in this graph , configuration should be pushed to the
// child policy and RPCs should get routed to that leaf cluster.
func (s) TestAggregatedCluster_CycleWithLeafNode(t *testing.T) {
	lbCfgCh, _, _, _ := registerWrappedClusterResolverPolicy(t)
	mgmtServer, nodeID, cc, _, _, _, _ := setupWithManagementServer(t)

	// Start a test service backend.
	server := stubserver.StartTestService(t, nil)
	t.Cleanup(server.Stop)

	const (
		clusterNameA = clusterName // cluster name in cds LB policy config
		clusterNameB = clusterName + "-B"
		clusterNameC = clusterName + "-C"
	)
	// Configure the management server with an aggregate cluster resource graph
	// that contains a cycle, but also contains a leaf cluster.
	resources := e2e.UpdateOptions{
		NodeID: nodeID,
		Clusters: []*v3clusterpb.Cluster{
			makeAggregateClusterResource(clusterNameA, []string{clusterNameB}),
			makeAggregateClusterResource(clusterNameB, []string{clusterNameA, clusterNameC}),
			e2e.DefaultCluster(clusterNameC, serviceName, e2e.SecurityLevelNone),
		},
		Endpoints:      []*v3endpointpb.ClusterLoadAssignment{e2e.DefaultEndpoint(serviceName, "localhost", []uint32{testutils.ParsePort(t, server.Address)})},
		SkipValidation: true,
	}
	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
	defer cancel()
	if err := mgmtServer.Update(ctx, resources); err != nil {
		t.Fatal(err)