k8s.io/kubernetes@v1.31.0-alpha.0.0.20240520171757-56147500dadc/test/utils/ktesting/withcontext.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  	"time"
    22  
    23  	"github.com/onsi/gomega"
    24  	"k8s.io/klog/v2"
    25  )
    26  
    27  // WithCancel sets up cancellation in a [TContext.Cleanup] callback and
    28  // constructs a new TContext where [TContext.Cancel] cancels only the new
    29  // context.
    30  func WithCancel(tCtx TContext) TContext {
    31  	ctx, cancel := context.WithCancelCause(tCtx)
    32  	tCtx.Cleanup(func() {
    33  		cancel(cleanupErr(tCtx.Name()))
    34  	})
    35  
    36  	return withContext{
    37  		TContext: tCtx,
    38  		Context:  ctx,
    39  		cancel: func(cause string) {
    40  			var cancelCause error
    41  			if cause != "" {
    42  				cancelCause = canceledError(cause)
    43  			}
    44  			cancel(cancelCause)
    45  		},
    46  	}
    47  }
    48  
    49  // WithTimeout sets up new context with a timeout. Canceling the timeout gets
    50  // registered in a [TContext.Cleanup] callback. [TContext.Cancel] cancels only
    51  // the new context. The cause is used as reason why the context is canceled
    52  // once the timeout is reached. It may be empty, in which case the usual
    53  // "context canceled" error is used.
    54  func WithTimeout(tCtx TContext, timeout time.Duration, timeoutCause string) TContext {
    55  	tCtx.Helper()
    56  	ctx, cancel := withTimeout(tCtx, tCtx.TB(), timeout, timeoutCause)
    57  
    58  	return withContext{
    59  		TContext: tCtx,
    60  		Context:  ctx,
    61  		cancel:   cancel,
    62  	}
    63  }
    64  
    65  // WithLogger constructs a new context with a different logger.
    66  func WithLogger(tCtx TContext, logger klog.Logger) TContext {
    67  	ctx := klog.NewContext(tCtx, logger)
    68  
    69  	return withContext{
    70  		TContext: tCtx,
    71  		Context:  ctx,
    72  		cancel:   tCtx.Cancel,
    73  	}
    74  }
    75  
    76  // withContext combines some TContext with a new [context.Context] derived
    77  // from it. Because both provide the [context.Context] interface, methods must
    78  // be defined which pick the newer one.
    79  type withContext struct {
    80  	TContext
    81  	context.Context
    82  
    83  	cancel func(cause string)
    84  }
    85  
    86  func (wCtx withContext) Cancel(cause string) {
    87  	wCtx.cancel(cause)
    88  }
    89  
    90  func (wCtx withContext) CleanupCtx(cb func(TContext)) {
    91  	wCtx.Helper()
    92  	cleanupCtx(wCtx, cb)
    93  }
    94  
    95  func (wCtx withContext) Expect(actual interface{}, extra ...interface{}) gomega.Assertion {
    96  	wCtx.Helper()
    97  	return expect(wCtx, actual, extra...)
    98  }
    99  
   100  func (wCtx withContext) ExpectNoError(err error, explain ...interface{}) {
   101  	wCtx.Helper()
   102  	expectNoError(wCtx, err, explain...)
   103  }
   104  
   105  func (wCtx withContext) Logger() klog.Logger {
   106  	return klog.FromContext(wCtx)
   107  }
   108  
   109  func (wCtx withContext) Deadline() (time.Time, bool) {
   110  	return wCtx.Context.Deadline()
   111  }
   112  
   113  func (wCtx withContext) Done() <-chan struct{} {
   114  	return wCtx.Context.Done()
   115  }
   116  
   117  func (wCtx withContext) Err() error {
   118  	return wCtx.Context.Err()
   119  }
   120  
   121  func (wCtx withContext) Value(key any) any {
   122  	return wCtx.Context.Value(key)
   123  }