xds/testing: export variables for testing (#4449)

The exported variables will be used by tests (to be added in a future
PR, in another package) that use these balancers as child balancer.

xds/internal/balancer/clusterimpl/balancer_test.go

@@ -66,7 +66,7 @@ func subConnFromPicker(p balancer.Picker) func() balancer.SubConn {
 }
 
 func init() {
-	newRandomWRR = testutils.NewTestWRR
+	NewRandomWRR = testutils.NewTestWRR
 }
 
 // TestDropByCategory verifies that the balancer correctly drops the picks, and

xds/internal/balancer/clusterimpl/picker.go

@@ -28,7 +28,9 @@ import (
 	"google.golang.org/grpc/xds/internal/client/load"
 )
 
-var newRandomWRR = wrr.NewRandom
+// NewRandomWRR is used when calculating drops. It's exported so that tests can
+// override it.
+var NewRandomWRR = wrr.NewRandom
 
 const million = 1000000
 
@@ -48,7 +50,7 @@ func gcd(a, b uint32) uint32 {
 }
 
 func newDropper(c DropConfig) *dropper {
-	w := newRandomWRR()
+	w := NewRandomWRR()
 	gcdv := gcd(c.RequestsPerMillion, million)
 	// Return true for RequestPerMillion, false for the rest.
 	w.Add(true, int64(c.RequestsPerMillion/gcdv))

xds/internal/balancer/priority/balancer_priority.go

@@ -28,8 +28,11 @@ import (
 )
 
 var (
-	errAllPrioritiesRemoved    = errors.New("no locality is provided, all priorities are removed")
-	defaultPriorityInitTimeout = 10 * time.Second
+	errAllPrioritiesRemoved = errors.New("no locality is provided, all priorities are removed")
+	// DefaultPriorityInitTimeout is the timeout after which if a priority is
+	// not READY, the next will be started. It's exported to be overridden by
+	// tests.
+	DefaultPriorityInitTimeout = 10 * time.Second
 )
 
 // syncPriority handles priority after a config update. It makes sure the
@@ -162,7 +165,7 @@ func (b *priorityBalancer) switchToChild(child *childBalancer, priority int) {
 		// to check the stopped boolean.
 		timerW := &timerWrapper{}
 		b.priorityInitTimer = timerW
-		timerW.timer = time.AfterFunc(defaultPriorityInitTimeout, func() {
+		timerW.timer = time.AfterFunc(DefaultPriorityInitTimeout, func() {
 			b.mu.Lock()
 			defer b.mu.Unlock()
 			if timerW.stopped {

xds/internal/balancer/priority/balancer_test.go

@@ -690,10 +690,10 @@ func (s) TestPriority_HigherReadyCloseAllLower(t *testing.T) {
 func (s) TestPriority_InitTimeout(t *testing.T) {
 	const testPriorityInitTimeout = time.Second
 	defer func() func() {
-		old := defaultPriorityInitTimeout
-		defaultPriorityInitTimeout = testPriorityInitTimeout
+		old := DefaultPriorityInitTimeout
+		DefaultPriorityInitTimeout = testPriorityInitTimeout
 		return func() {
-			defaultPriorityInitTimeout = old
+			DefaultPriorityInitTimeout = old
 		}
 	}()()
 
@@ -760,10 +760,10 @@ func (s) TestPriority_InitTimeout(t *testing.T) {
 func (s) TestPriority_RemovesAllPriorities(t *testing.T) {
 	const testPriorityInitTimeout = time.Second
 	defer func() func() {
-		old := defaultPriorityInitTimeout
-		defaultPriorityInitTimeout = testPriorityInitTimeout
+		old := DefaultPriorityInitTimeout
+		DefaultPriorityInitTimeout = testPriorityInitTimeout
 		return func() {
-			defaultPriorityInitTimeout = old
+			DefaultPriorityInitTimeout = old
 		}
 	}()()
 
@@ -1030,9 +1030,9 @@ func (s) TestPriority_HighPriorityNoEndpoints(t *testing.T) {
 func (s) TestPriority_FirstPriorityUnavailable(t *testing.T) {
 	const testPriorityInitTimeout = time.Second
 	defer func(t time.Duration) {
-		defaultPriorityInitTimeout = t
-	}(defaultPriorityInitTimeout)
-	defaultPriorityInitTimeout = testPriorityInitTimeout
+		DefaultPriorityInitTimeout = t
+	}(DefaultPriorityInitTimeout)
+	DefaultPriorityInitTimeout = testPriorityInitTimeout
 
 	cc := testutils.NewTestClientConn(t)
 	bb := balancer.Get(Name)

xds/internal/balancer/weightedtarget/weightedtarget.go

@@ -37,9 +37,9 @@ import (
 // Name is the name of the weighted_target balancer.
 const Name = "weighted_target_experimental"
 
-// newRandomWRR is the WRR constructor used to pick sub-pickers from
+// NewRandomWRR is the WRR constructor used to pick sub-pickers from
 // sub-balancers. It's to be modified in tests.
-var newRandomWRR = wrr.NewRandom
+var NewRandomWRR = wrr.NewRandom
 
 func init() {
 	balancer.Register(&weightedTargetBB{})
@@ -50,7 +50,7 @@ type weightedTargetBB struct{}
 func (wt *weightedTargetBB) Build(cc balancer.ClientConn, bOpts balancer.BuildOptions) balancer.Balancer {
 	b := &weightedTargetBalancer{}
 	b.logger = prefixLogger(b)
-	b.stateAggregator = weightedaggregator.New(cc, b.logger, newRandomWRR)
+	b.stateAggregator = weightedaggregator.New(cc, b.logger, NewRandomWRR)
 	b.stateAggregator.Start()
 	b.bg = balancergroup.New(cc, bOpts, b.stateAggregator, nil, b.logger)
 	b.bg.Start()