Repository: grpc/grpc-go
File: balancer/randomsubsetting/randomsubsetting_test.go (Go)

Code before:
			  "random_subsetting": {
				"subset_size": %d,
				"child_policy": [{"round_robin": {}}]
			  }
			}
		  ]
		}`, subsetSize)

		sc := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(string(jsonConfig))
		mr.InitialState(resolver.State{
			Addresses:     addresses,
			ServiceConfig: sc,
		})

		ccs[i], err = grpc.Dial(mr.Scheme()+":///", grpc.WithResolvers(mr), grpc.WithTransportCredentials(insecure.NewCredentials()))
		if err != nil {
			t.Fatalf("grpc.Dial() failed: %v", err)
		}
		clients[i] = testgrpc.NewTestServiceClient(ccs[i])
	}

	cancel := func() {
		for _, cc := range ccs {
			cc.Close()
		}
	}
	return clients, cancel
}

func checkRoundRobinRPCs(ctx context.Context, t *testing.T, clients []testgrpc.TestServiceClient, subsetSize int, maxDiff int) {
	clientsPerBackend := map[string]map[int]struct{}{}

	for clientIdx, client := range clients {
		// make sure that every client send exactly 1 request to each server in its subset
		for i := 0; i < subsetSize; i++ {
			var peer peer.Peer
			_, err := client.EmptyCall(ctx, &testpb.Empty{}, grpc.Peer(&peer))
			if err != nil {
				t.Fatalf("failed to call server: %v", err)
			}
			if peer.Addr != nil {
				if m, ok := clientsPerBackend[peer.Addr.String()]; !ok {
					clientsPerBackend[peer.Addr.String()] = map[int]struct{}{clientIdx: {}}
				} else if _, ok := m[clientIdx]; !ok {
					m[clientIdx] = struct{}{}
				} else {
					// The backend receives a second request from the same client. This could happen if the client have 1 backend in READY
					// state while the other  are CONNECTING. In this case round_robbin will pick the same address twice.
					// We are going to retry after short timeout.
					time.Sleep(10 * time.Microsecond)
					i--

Review comment (suggestion):
You could use `t.Cleanup` here as well and avoid having to return a cancel function that the test needs to invoke.

Code after:
			}
		})
	}
}

func (s) TestCalculateSubset_EndpointsRetainHashValues(t *testing.T) {
	endpoints := makeEndpoints(10)
	const subsetSize = 5
	// The subset is deterministic based on the hash, so we can hardcode
	// the expected output.
	want := []resolver.Endpoint{
		{Addresses: []resolver.Address{{Addr: "endpoint-6"}}},
		{Addresses: []resolver.Address{{Addr: "endpoint-0"}}},
		{Addresses: []resolver.Address{{Addr: "endpoint-1"}}},
		{Addresses: []resolver.Address{{Addr: "endpoint-7"}}},
		{Addresses: []resolver.Address{{Addr: "endpoint-3"}}},
	}

	b := &subsettingBalancer{
		cfg:        &lbConfig{SubsetSize: subsetSize},
		hashSeed:   0,
		hashDigest: xxhash.New(),
	}
	for range 10 {
		got := b.calculateSubset(endpoints)
		if diff := cmp.Diff(want, got); diff != "" {
			t.Fatalf("calculateSubset() returned diff (-want +got):\n%s", diff)
		}
	}
}

func (s) TestSubsettingBalancer_DeterministicSubset(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
	defer cancel()

	// Register the stub balancer builder, which will be used as the child
	// policy in the random_subsetting balancer.
	updateCh := make(chan balancer.ClientConnState, 1)
	stub.Register("stub-child-balancer", stub.BalancerFuncs{
		UpdateClientConnState: func(_ *stub.BalancerData, s balancer.ClientConnState) error {
			select {
			case <-ctx.Done():
			case updateCh <- s:
			}
			return nil
		},
	})

	// Create a random_subsetting balancer.
	tcc := testutils.NewBalancerClientConn(t)
	rsb := balancer.Get(Name).Build(tcc, balancer.BuildOptions{})