Repository: grpc/grpc-go
File: balancer/rls/balancer_test.go (Go)

Code before:
    }
  ]
}`, topLevelBalancerName, rlsServer.Address, childPolicyName)
	sc = internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(scJSON)
	r.UpdateState(resolver.State{ServiceConfig: sc})

	// Wait for the clientconn update to be processed by the RLS LB policy.
	select {
	case <-ctx.Done():
	case <-clientConnUpdateDone:
	}

	// Even though the child policies used in this test make multiple calls to
	// UpdateState as part of handling their configs, we expect the RLS policy
	// to inhibit picker updates during this time frame, and send a single
	// picker once the config update is completely handled.
	states1 := ccWrapper.getStates()
	if len(states1) != len(states0)+1 {
		t.Fatalf("more than one state update seen. before %v, after %v", states0, states1)
	}
}

// verifyStateTransitions verifies that the given number of state updates are
// received on the channel and they match the expected sequence. Adjacent
// duplicates of the waiting state are ignored.
func verifyStateTransitions(ctx context.Context, t *testing.T, stateCh <-chan balancer.State, wantNum int, wantStates ...connectivity.State) {
	t.Helper()

	currIdx := 0 // Index of the state we are currently waiting for.
	for i := 0; i < wantNum; i++ {
		select {
		case state := <-stateCh:
			wantState := wantStates[currIdx]
			if state.ConnectivityState == wantState {
				if currIdx != len(wantStates)-1 {
					currIdx++
				}
				continue
			}

			// Ignore duplicates of the *previous* state we already verified.
			if currIdx > 0 && state.ConnectivityState == wantStates[currIdx-1] {
				continue
			}
			t.Fatalf("Got unexpected balancer state update: %v, want %v", state.ConnectivityState, wantState)
		case <-ctx.Done():
			t.Fatalf("Timeout waiting for balancer state update %v", wantStates[currIdx])
		}
	}
}


Review comment (refactor):
As I understand it, this helper performs two jobs:
1. Waiting for `wantNum` state updates.
2. De-duplicating consecutive state updates.

This combination makes it a little hard to read and potentially reuse in other tests. In my opinion, it would be better if this only handled waiting for `wantNum` and returned the states as a slice to the caller. The caller can then do its own custom validation.

This would make the code easier to understand since the waiting and the de-duplication would happen serially. We could also use [slices.Compact](https://pkg.go.dev/slices#Compact) or [CompactFunc](https://pkg.go.dev/slices#CompactFunc) for the de-duplication logic.

Code after:
	// update on the parent. Hence, this update is required to force the
	// scenario which we are interesting in testing here, i.e child policies get
	// config updates as part of the parent policy getting its config update.
	scJSON = fmt.Sprintf(`
{
  "loadBalancingConfig": [
    {
      "%s": {
		"routeLookupConfig": {
			"grpcKeybuilders": [{
				"names": [
					{"service": "grpc.testing.TestService"},
					{"service": "grpc.health.v1.Health"}
				]
			}],
			"lookupService": "%s",
			"cacheSizeBytes": 1000
		},
		"childPolicy": [{"%s": {}}],
		"childPolicyConfigTargetFieldName": "Backend"
      }
    }
  ]
}`, topLevelBalancerName, rlsServer.Address, childPolicyName)
	sc = internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(scJSON)
	r.UpdateState(resolver.State{ServiceConfig: sc})

	// Wait for the clientconn update to be processed by the RLS LB policy.
	select {
	case <-ctx.Done():
	case <-clientConnUpdateDone:
	}

	// Even though the child policies used in this test make multiple calls to
	// UpdateState as part of handling their configs, we expect the RLS policy
	// to inhibit picker updates during this time frame, and send a single
	// picker once the config update is completely handled.
	states1 := ccWrapper.getStates()
	if len(states1) != len(states0)+1 {
		t.Fatalf("more than one state update seen. before %v, after %v", states0, states1)
	}
}

// waitForStateTransitions waits for the given number of state updates on the
// channel and returns the sequence of connectivity states received.
func waitForStateTransitions(ctx context.Context, t *testing.T, stateCh <-chan balancer.State, wantNum int) []connectivity.State {
	t.Helper()

	var gotStates []connectivity.State
	for i := 0; i < wantNum; i++ {
		select {