Repository: grpc/grpc-go
File: internal/xds/balancer/clusterimpl/tests/balancer_test.go (Go)

Code before:
	defer cc.Close()

	client := testgrpc.NewTestServiceClient(cc)
	if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
		t.Fatalf("client.EmptyCall() failed: %v", err)
	}

	// Register stub pickfirst LB policy so that we can catch config changes.
	pfBuilder := balancer.Get(pickfirst.Name)
	internal.BalancerUnregister(pfBuilder.Name())
	lbCfgCh := make(chan serviceconfig.LoadBalancingConfig, 1)
	updatedChildPolicy := atomic.Pointer[string]{}
	stub.Register(pfBuilder.Name(), stub.BalancerFuncs{
		ParseConfig: func(lbCfg json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
			return pfBuilder.(balancer.ConfigParser).ParseConfig(lbCfg)
		},
		Init: func(bd *stub.BalancerData) {
			bd.ChildBalancer = pfBuilder.Build(bd.ClientConn, bd.BuildOptions)
		},
		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
			name := pfBuilder.Name()
			updatedChildPolicy.Store(&name)
			select {
			case lbCfgCh <- ccs.BalancerConfig:
			case <-ctx.Done():
				t.Fatalf("Timeout while waiting for BalancerConfig")
			}
			return bd.ChildBalancer.UpdateClientConnState(ccs)
		},
		Close: func(bd *stub.BalancerData) {
			bd.ChildBalancer.Close()
		},
	})
	defer balancer.Register(pfBuilder)

	// Now update the cluster to use "pick_first" as the endpoint picking policy.
	resources.Clusters[0].LoadBalancingPolicy = &v3clusterpb.LoadBalancingPolicy{
		Policies: []*v3clusterpb.LoadBalancingPolicy_Policy{{
			TypedExtensionConfig: &v3corepb.TypedExtensionConfig{
				TypedConfig: testutils.MarshalAny(t, &v3pickfirstpb.PickFirst{}),
			},
		}},
	}
	if err := mgmtServer.Update(ctx, resources); err != nil {
		t.Fatal(err)
	}

	select {
	case <-ctx.Done():
		t.Fatalf("Timeout waiting for pickfirst child policy config")
	case <-lbCfgCh:

Review comment (bug):
Hmm .. I would consider changing this to a `t.Error` because this piece of code would generally be executed from another goroutine and not from the main test goroutine. Only calls to `t.Fatal` from the main test goroutine lead to immediately test failures. Calls to `t.Fatal` from non main test goroutine would behave the same as .t.Error`

Code after:
	cc, err := grpc.NewClient(fmt.Sprintf("xds:///%s", serviceName), grpc.WithResolvers(resolverBuilder), grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		t.Fatalf("Failed to create client: %v", err)
	}
	defer cc.Close()

	client := testgrpc.NewTestServiceClient(cc)
	if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
		t.Fatalf("client.EmptyCall() failed: %v", err)
	}

	// Register stub pickfirst LB policy so that we can catch config changes.
	pfBuilder := balancer.Get(pickfirst.Name)
	internal.BalancerUnregister(pfBuilder.Name())
	lbCfgCh := make(chan serviceconfig.LoadBalancingConfig, 1)
	var updatedChildPolicy atomic.Pointer[string]
	stub.Register(pfBuilder.Name(), stub.BalancerFuncs{
		ParseConfig: func(lbCfg json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
			return pfBuilder.(balancer.ConfigParser).ParseConfig(lbCfg)
		},
		Init: func(bd *stub.BalancerData) {
			bd.ChildBalancer = pfBuilder.Build(bd.ClientConn, bd.BuildOptions)
		},
		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
			name := pfBuilder.Name()
			updatedChildPolicy.Store(&name)
			select {
			case lbCfgCh <- ccs.BalancerConfig:
			case <-ctx.Done():
				t.Error("Timed out while waiting for BalancerConfig, context deadline exceeded")
			}
			return bd.ChildBalancer.UpdateClientConnState(ccs)
		},
		Close: func(bd *stub.BalancerData) {
			bd.ChildBalancer.Close()
		},
	})
	defer balancer.Register(pfBuilder)

	// Now update the cluster to use "pick_first" as the endpoint picking policy.
	resources.Clusters[0].LoadBalancingPolicy = &v3clusterpb.LoadBalancingPolicy{
		Policies: []*v3clusterpb.LoadBalancingPolicy_Policy{{
			TypedExtensionConfig: &v3corepb.TypedExtensionConfig{
				TypedConfig: testutils.MarshalAny(t, &v3pickfirstpb.PickFirst{}),
			},
		}},
	}
	if err := mgmtServer.Update(ctx, resources); err != nil {
		t.Fatal(err)
	}