grpc: Fix flaky picker_wrapper tests

[arjan-bal]

Fixes: #7521

Background

The tests call picker_wrapper.pick() in 5 go routines. The go routines are blocked on either the context getting cancelled or the picker being updated.

grpc-go/picker_wrapper.go

Lines 120 to 134 in cfd14ba

	select {
	case <-ctx.Done():
	var errStr string
	if lastPickErr != nil {
	errStr = "latest balancer error: " + lastPickErr.Error()
	} else {
	errStr = fmt.Sprintf("received context error while waiting for new LB policy update: %s", ctx.Err().Error())
	}
	switch ctx.Err() {
	case context.DeadlineExceeded:
	return nil, balancer.PickResult{}, status.Error(codes.DeadlineExceeded, errStr)
	case context.Canceled:
	return nil, balancer.PickResult{}, status.Error(codes.Canceled, errStr)
	}
	case <-ch:

The context is cancelled through a deferred statement when the test function exits. The test function updates the picker in the last line before exiting:

grpc-go/picker_wrapper_test.go

Line 162 in cfd14ba

bp.updatePicker(&testingPicker{sc: testSC, maxCalled: goroutineCount})

When the picker go routines see the context getting cancelled before the picker update, they log an error causing the test to fail.

Repro

Add a time.Sleep(100 * time.Milliseconds) just before waiting on the picker update / context cancellation here:

grpc-go/picker_wrapper.go

Lines 116 to 120 in cfd14ba

	if ch == pg.blockingCh {
	// This could happen when either:
	// - pw.picker is nil (the previous if condition), or
	// - we have already called pick on the current picker.
	select {

This will cause the test to fail with the following log picker_wrapper_test.go:153: bp.pick returned non-nil error: rpc error: code = Canceled desc = received context error while waiting for new LB policy update: context canceled

Fix

Wait for all the go routines to finish before returning from the test function and cancelling the context.

RELEASE NOTES: None

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 81.70%. Comparing base (b45fc41) to head (5fa62ae).
Report is 15 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #7560      +/-   ##
==========================================
- Coverage   82.07%   81.70%   -0.38%     
==========================================
  Files         360      361       +1     
  Lines       27533    27816     +283     
==========================================
+ Hits        22599    22726     +127     
- Misses       3759     3876     +117     
- Partials     1175     1214      +39     

see 51 files with indirect coverage changes

[easwars]

Good find. LGTM modulo the minor nit.

[easwars]

Nit: While you are here, could you please improve these error messages to say something like:

t.Errorf("bp.pick returned transport: %v, error: %v, want transport: %v, error: nil", tr, err, testT)

[arjan-bal]

Changed the error messages for the 4 tests.