Repository: grpc/grpc-go
File: test/balancer_switching_test.go (Go)

Code before:
	// lead to the channel using the default LB policy which is pick_first. The
	// list of addresses pushed as part of this update is different from the one
	// returned by the "grpclb" balancer. So, we should see RPCs going to the
	// newly configured backends, as part of the balancer switch.
	emptyConfig := parseServiceConfig(t, r, `{}`)
	r.UpdateState(resolver.State{Addresses: addrs[1:], ServiceConfig: emptyConfig})
	if err := pfutil.CheckRPCsToBackend(ctx, cc, addrs[1]); err != nil {
		t.Fatal(err)
	}
}

// TestBalancerSwitch_pickFirstToGRPCLB tests the scenario where the channel
// starts off with "pick_first", switches to "grpclb" and back.
func (s) TestBalancerSwitch_pickFirstToGRPCLB(t *testing.T) {
	backends, lbServer, cleanup := setupBackendsAndFakeGRPCLB(t)
	defer cleanup()

	addrs := stubBackendsToResolverAddrs(backends)
	r := manual.NewBuilderWithScheme("whatever")
	target := fmt.Sprintf("%s:///%s", r.Scheme(), loadBalancedServiceName)
	cc, err := grpc.NewClient(target, grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithResolvers(r))
	if err != nil {
		t.Fatalf("grpc.NewClient() failed: %v", err)
	}
	defer cc.Close()
	cc.Connect()
	// Push a resolver update containing no grpclb server address. This should
	// lead to the channel using the default LB policy which is pick_first.
	r.UpdateState(resolver.State{Addresses: addrs[1:]})
	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
	defer cancel()
	if err := pfutil.CheckRPCsToBackend(ctx, cc, addrs[1]); err != nil {
		t.Fatal(err)
	}

	// Push a resolver update with no service config and a single address pointing
	// to the grpclb server we created above. This will cause the channel to
	// switch to the "grpclb" balancer, which returns a single backend address.
	grpclbConfig := parseServiceConfig(t, r, `{"loadBalancingPolicy": "grpclb"}`)
	state := resolver.State{ServiceConfig: grpclbConfig}
	r.UpdateState(grpclbstate.Set(state, &grpclbstate.State{BalancerAddresses: []resolver.Address{{Addr: lbServer.Address()}}}))
	client := testgrpc.NewTestServiceClient(cc)
	if err := rrutil.CheckRoundRobinRPCs(ctx, client, addrs[:1]); err != nil {
		t.Fatal(err)
	}

	// Push a resolver update containing a non-existent grpclb server address.
	// This should not lead to a balancer switch.
	r.UpdateState(grpclbstate.Set(state, &grpclbstate.State{BalancerAddresses: []resolver.Address{{Addr: "nonExistentServer"}}}))
	if err := rrutil.CheckRoundRobinRPCs(ctx, client, addrs[:1]); err != nil {
		t.Fatal(err)

Review comment (suggestion):
You can avoid using `cc.Connect` here by replacing the `r.UpdateState` below with `r.InitialState`. Please check other places where the same change can be made.

Code after:

	// Push a resolver update containing no grpclb server address. This should
	// lead to the channel using the default LB policy which is pick_first. The
	// list of addresses pushed as part of this update is different from the one
	// returned by the "grpclb" balancer. So, we should see RPCs going to the
	// newly configured backends, as part of the balancer switch.
	emptyConfig := parseServiceConfig(t, r, `{}`)
	r.UpdateState(resolver.State{Addresses: addrs[1:], ServiceConfig: emptyConfig})
	if err := pfutil.CheckRPCsToBackend(ctx, cc, addrs[1]); err != nil {
		t.Fatal(err)
	}
}

// TestBalancerSwitch_pickFirstToGRPCLB tests the scenario where the channel
// starts off with "pick_first", switches to "grpclb" and back.
func (s) TestBalancerSwitch_pickFirstToGRPCLB(t *testing.T) {
	backends, lbServer, cleanup := setupBackendsAndFakeGRPCLB(t)
	defer cleanup()

	addrs := stubBackendsToResolverAddrs(backends)
	r := manual.NewBuilderWithScheme("whatever")
	target := fmt.Sprintf("%s:///%s", r.Scheme(), loadBalancedServiceName)
	cc, err := grpc.NewClient(target, grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithResolvers(r))
	if err != nil {
		t.Fatalf("grpc.NewClient() failed: %v", err)
	}
	defer cc.Close()

	// Set an empty initial resolver state. This should lead to the channel
	// using the default LB policy which is pick_first.
	r.InitialState(resolver.State{Addresses: addrs[1:]})

	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
	defer cancel()
	if err := pfutil.CheckRPCsToBackend(ctx, cc, addrs[1]); err != nil {
		t.Fatal(err)
	}

	// Push a resolver update with no service config and a single address pointing
	// to the grpclb server we created above. This will cause the channel to
	// switch to the "grpclb" balancer, which returns a single backend address.
	grpclbConfig := parseServiceConfig(t, r, grpclbServiceConfig)
	state := resolver.State{ServiceConfig: grpclbConfig}
	r.UpdateState(grpclbstate.Set(state, &grpclbstate.State{BalancerAddresses: []resolver.Address{{Addr: lbServer.Address()}}}))
	client := testgrpc.NewTestServiceClient(cc)
	if err := rrutil.CheckRoundRobinRPCs(ctx, client, addrs[:1]); err != nil {
		t.Fatal(err)
	}

	// Push a resolver update containing a non-existent grpclb server address.
	// This should not lead to a balancer switch.