extract context creation from some more loops

Author: hasson82

picker_wrapper_test.go

@@ -99,12 +99,11 @@ func (s) TestBlockingPickNoSubAvailable(t *testing.T) {
 	bp := newPickerWrapper(nil)
 	var finishedCount uint64
 	bp.updatePicker(&testingPicker{err: balancer.ErrNoSubConnAvailable, maxCalled: goroutineCount})
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
 	// All goroutines should block because picker returns no subConn available.
 	for i := goroutineCount; i > 0; i-- {
 		go func() {
-			ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
-			defer cancel()
-
 			if tr, _, err := bp.pick(ctx, true, balancer.PickInfo{}); err != nil || tr != testT {
 				t.Errorf("bp.pick returned non-nil error: %v", err)
 			}
@@ -122,13 +121,12 @@ func (s) TestBlockingPickTransientWaitforready(t *testing.T) {
 	bp := newPickerWrapper(nil)
 	bp.updatePicker(&testingPicker{err: balancer.ErrTransientFailure, maxCalled: goroutineCount})
 	var finishedCount uint64
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
 	// All goroutines should block because picker returns transientFailure and
 	// picks are not failfast.
 	for i := goroutineCount; i > 0; i-- {
 		go func() {
-			ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
-			defer cancel()
-
 			if tr, _, err := bp.pick(ctx, false, balancer.PickInfo{}); err != nil || tr != testT {
 				t.Errorf("bp.pick returned non-nil error: %v", err)
 			}
@@ -146,12 +144,11 @@ func (s) TestBlockingPickSCNotReady(t *testing.T) {
 	bp := newPickerWrapper(nil)
 	bp.updatePicker(&testingPicker{sc: testSCNotReady, maxCalled: goroutineCount})
 	var finishedCount uint64
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
 	// All goroutines should block because subConn is not ready.
 	for i := goroutineCount; i > 0; i-- {
 		go func() {
-			ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
-			defer cancel()
-
 			if tr, _, err := bp.pick(ctx, true, balancer.PickInfo{}); err != nil || tr != testT {
 				t.Errorf("bp.pick returned non-nil error: %v", err)
 			}

server_test.go

@@ -192,10 +192,9 @@ func BenchmarkChainUnaryInterceptor(b *testing.B) {
 			s := NewServer(ChainUnaryInterceptor(interceptors...))
 			b.ReportAllocs()
 			b.ResetTimer()
+			ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+			defer cancel()
 			for i := 0; i < b.N; i++ {
-				ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
-				defer cancel()
-
 				if _, err := s.opts.unaryInt(ctx, nil, nil,
 					func(ctx context.Context, req any) (any, error) {
 						return nil, nil