priority: add a test helper to override the init timeout in tests (#8704)

This is an initial cleanup before the actual work required for
https://github.com/grpc/grpc-go/issues/8516

The way the init timeout is currently overridden in existing tests is
unnecessary complicated and hard to read. This simplifies things and
ensures that the previous pattern is not followed when new tests are
written.

RELEASE NOTES: none

diff --git a/internal/xds/balancer/priority/balancer_test.go b/internal/xds/balancer/priority/balancer_test.go
index 3aac85ce54f8f4b1d8f99444e889a4270520ddaf..dbded5317af910755d6a238ff24afc885cfc9bfa 100644 (file)
--- a/internal/xds/balancer/priority/balancer_test.go
+++ b/internal/xds/balancer/priority/balancer_test.go
@@ -74,6 +74,12 @@ func init() {
        balancer.Register(&anotherRR{Builder: balancer.Get(roundrobin.Name)})
 }
 
+func overrideInitTimeout(t *testing.T, val time.Duration) {
+       orig := DefaultPriorityInitTimeout
+       DefaultPriorityInitTimeout = val
+       t.Cleanup(func() { DefaultPriorityInitTimeout = orig })
+}
+
 // When a high priority is ready, adding/removing lower locality doesn't cause
 // changes.
 //
@@ -674,13 +680,7 @@ func (s) TestPriority_InitTimeout(t *testing.T) {
        defer cancel()
 
        const testPriorityInitTimeout = 200 * time.Millisecond
-       defer func() func() {
-               old := DefaultPriorityInitTimeout
-               DefaultPriorityInitTimeout = testPriorityInitTimeout
-               return func() {
-                       DefaultPriorityInitTimeout = old
-               }
-       }()()
+       overrideInitTimeout(t, testPriorityInitTimeout)
 
        cc := testutils.NewBalancerClientConn(t)
        bb := balancer.Get(Name)
@@ -745,13 +745,7 @@ func (s) TestPriority_RemovesAllPriorities(t *testing.T) {
        defer cancel()
 
        const testPriorityInitTimeout = 200 * time.Millisecond
-       defer func() func() {
-               old := DefaultPriorityInitTimeout
-               DefaultPriorityInitTimeout = testPriorityInitTimeout
-               return func() {
-                       DefaultPriorityInitTimeout = old
-               }
-       }()()
+       overrideInitTimeout(t, testPriorityInitTimeout)
 
        cc := testutils.NewBalancerClientConn(t)
        bb := balancer.Get(Name)
@@ -1011,10 +1005,7 @@ func (s) TestPriority_HighPriorityNoEndpoints(t *testing.T) {
 // Test the case where the first and only priority is removed.
 func (s) TestPriority_FirstPriorityUnavailable(t *testing.T) {
        const testPriorityInitTimeout = 200 * time.Millisecond
-       defer func(t time.Duration) {
-               DefaultPriorityInitTimeout = t
-       }(DefaultPriorityInitTimeout)
-       DefaultPriorityInitTimeout = testPriorityInitTimeout
+       overrideInitTimeout(t, testPriorityInitTimeout)
 
        cc := testutils.NewBalancerClientConn(t)
        bb := balancer.Get(Name)