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

Code before:

	// Push another resolver update with a service config specifying "pick_first".
	r.UpdateState(resolver.State{
		Addresses:     addrs,
		ServiceConfig: parseServiceConfig(t, r, pickFirstServiceConfig),
	})
	if err := pfutil.CheckRPCsToBackend(ctx, cc, addrs[0]); err != nil {
		t.Fatal(err)
	}
}

// TestBalancerSwitch_grpclbToPickFirst tests the scenario where the channel
// starts off "grpclb", switches to "pick_first" and back.
func (s) TestBalancerSwitch_grpclbToPickFirst(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)
	}
	cc.Connect()
	defer cc.Close()

	// Set an initial resolver with a GRPCLB 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)
	r.InitialState(resolver.State{
		ServiceConfig: grpclbConfig,
		Addresses:     []resolver.Address{{Addr: lbServer.Address()}},
	})

	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
	defer cancel()
	client := testgrpc.NewTestServiceClient(cc)
	if err := rrutil.CheckRoundRobinRPCs(ctx, client, addrs[0: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.
	const nonExistentServer = "non-existent-grpclb-server-address"
	r.UpdateState(resolver.State{
		ServiceConfig: grpclbConfig,
		Addresses:     []resolver.Address{{Addr: nonExistentServer}},

Review comment (bug):
Calling `Connect()` here seems wrong. `InitialState` should be called before calling `Connect()`. In this case we probably don't even need to call `Connect` as `rrutil.CheckRoundRobinRPCs` makes RPCs which make the channel exit idle state automatically.

Code after:
	r.UpdateState(resolver.State{
		Addresses:     addrs,
		ServiceConfig: parseServiceConfig(t, r, pickFirstServiceConfig),
	})
	if err := pfutil.CheckRPCsToBackend(ctx, cc, addrs[0]); err != nil {
		t.Fatal(err)
	}
}

// TestBalancerSwitch_grpclbToPickFirst tests the scenario where the channel
// starts off "grpclb", switches to "pick_first" and back.
func (s) TestBalancerSwitch_grpclbToPickFirst(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 with a GRPCLB 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()}}}))
	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
	defer cancel()
	client := testgrpc.NewTestServiceClient(cc)
	if err := rrutil.CheckRoundRobinRPCs(ctx, client, addrs[0: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.
	const nonExistentServer = "non-existent-grpclb-server-address"
	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)
	}

	// 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