k8s.io/kubernetes@v1.31.0-alpha.0.0.20240520171757-56147500dadc/test/utils/ktesting/contexthelper.go (about)

     1  /*
     2  Copyright 2023 The Kubernetes Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package ktesting
    18  
    19  import (
    20  	"context"
    21  	"fmt"
    22  	"time"
    23  )
    24  
    25  // cleanupErr creates a cause when canceling a context because the test has completed.
    26  // It is a context.Canceled error.
    27  func cleanupErr(testName string) error {
    28  	return canceledError(fmt.Sprintf("test %s is cleaning up", testName))
    29  }
    30  
    31  type canceledError string
    32  
    33  func (c canceledError) Error() string { return string(c) }
    34  
    35  func (c canceledError) Is(target error) bool {
    36  	return target == context.Canceled
    37  }
    38  
    39  // withTimeout corresponds to [context.WithTimeout]. In contrast to
    40  // [context.WithTimeout], it automatically cancels during test cleanup, provides
    41  // the given cause when the deadline is reached, and its cancel function
    42  // requires a cause.
    43  func withTimeout(ctx context.Context, tb TB, timeout time.Duration, timeoutCause string) (context.Context, func(cause string)) {
    44  	tb.Helper()
    45  
    46  	now := time.Now()
    47  
    48  	cancelCtx, cancel := context.WithCancelCause(ctx)
    49  	after := time.NewTimer(timeout)
    50  	stopCtx, stop := context.WithCancel(ctx) // Only used internally, doesn't need a cause.
    51  	tb.Cleanup(func() {
    52  		cancel(cleanupErr(tb.Name()))
    53  		stop()
    54  	})
    55  	go func() {
    56  		select {
    57  		case <-stopCtx.Done():
    58  			after.Stop()
    59  			// No need to set a cause here. The cause or error of
    60  			// the parent context will be used.
    61  		case <-after.C:
    62  			cancel(canceledError(timeoutCause))
    63  		}
    64  	}()
    65  
    66  	// Determine which deadline is sooner: ours or that of our parent.
    67  	deadline := now.Add(timeout)
    68  	if parentDeadline, ok := ctx.Deadline(); ok {
    69  		if deadline.After(parentDeadline) {
    70  			deadline = parentDeadline
    71  		}
    72  	}
    73  
    74  	// We always have a deadline.
    75  	return deadlineContext{Context: cancelCtx, deadline: deadline}, func(cause string) {
    76  		var cancelCause error
    77  		if cause != "" {
    78  			cancelCause = canceledError(cause)
    79  		}
    80  		cancel(cancelCause)
    81  	}
    82  }
    83  
    84  type deadlineContext struct {
    85  	context.Context
    86  	deadline time.Time
    87  }
    88  
    89  func (d deadlineContext) Deadline() (time.Time, bool) {
    90  	return d.deadline, true
    91  }