change all usage of context in server_test.go to context.WithTimeout

Author: hasson82

server_test.go

@@ -153,15 +153,21 @@ func (s) TestRetryChainedInterceptor(t *testing.T) {
 	handler := func(ctx context.Context, req any) (any, error) {
 		return nil, nil
 	}
-	ii(context.Background(), nil, nil, handler)
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	ii(ctx, nil, nil, handler)
 	if !cmp.Equal(records, []int{1, 2, 3, 2, 3}) {
 		t.Fatalf("retry failed on chained interceptors: %v", records)
 	}
 }
 
 func (s) TestStreamContext(t *testing.T) {
 	expectedStream := &transport.Stream{}
-	ctx := NewContextWithServerTransportStream(context.Background(), expectedStream)
+	ctx, cancel := context.WithTimeout(NewContextWithServerTransportStream(context.Background(), expectedStream), defaultTestTimeout)
+	defer cancel()
+
 	s := ServerTransportStreamFromContext(ctx)
 	stream, ok := s.(*transport.Stream)
 	if !ok || expectedStream != stream {
@@ -186,7 +192,10 @@ func BenchmarkChainUnaryInterceptor(b *testing.B) {
 			b.ReportAllocs()
 			b.ResetTimer()
 			for i := 0; i < b.N; i++ {
-				if _, err := s.opts.unaryInt(context.Background(), nil, nil,
+				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
 					},