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 "flag" 22 "fmt" 23 "time" 24 25 "github.com/onsi/gomega" 26 apiextensions "k8s.io/apiextensions-apiserver/pkg/client/clientset/clientset" 27 "k8s.io/client-go/dynamic" 28 clientset "k8s.io/client-go/kubernetes" 29 "k8s.io/client-go/rest" 30 "k8s.io/client-go/restmapper" 31 "k8s.io/klog/v2" 32 "k8s.io/klog/v2/ktesting" 33 "k8s.io/kubernetes/test/utils/format" 34 "k8s.io/kubernetes/test/utils/ktesting/initoption" 35 "k8s.io/kubernetes/test/utils/ktesting/internal" 36 ) 37 38 // Underlier is the additional interface implemented by the per-test LogSink 39 // behind [TContext.Logger]. 40 type Underlier = ktesting.Underlier 41 42 // CleanupGracePeriod is the time that a [TContext] gets canceled before the 43 // deadline of its underlying test suite (usually determined via "go test 44 // -timeout"). This gives the running test(s) time to fail with an informative 45 // timeout error. After that, all cleanup callbacks then have the remaining 46 // time to complete before the test binary is killed. 47 // 48 // For this to work, each blocking calls in a test must respect the 49 // cancellation of the [TContext]. 50 // 51 // When using Ginkgo to manage the test suite and running tests, the 52 // CleanupGracePeriod is ignored because Ginkgo itself manages timeouts. 53 const CleanupGracePeriod = 5 * time.Second 54 55 // TContext combines [context.Context], [TB] and some additional 56 // methods. Log output is associated with the current test. Errors ([Error], 57 // [Errorf]) are recorded with "ERROR" as prefix, fatal errors ([Fatal], 58 // [Fatalf]) with "FATAL ERROR". 59 // 60 // TContext provides features offered by Ginkgo also when using normal Go [testing]: 61 // - The context contains a deadline that expires soon enough before 62 // the overall timeout that cleanup code can still run. 63 // - Cleanup callbacks can get their own, separate contexts when 64 // registered via [CleanupCtx]. 65 // - CTRL-C aborts, prints a progress report, and then cleans up 66 // before terminating. 67 // - SIGUSR1 prints a progress report without aborting. 68 // 69 // Progress reporting is more informative when doing polling with 70 // [gomega.Eventually] and [gomega.Consistently]. Without that, it 71 // can only report which tests are active. 72 type TContext interface { 73 context.Context 74 TB 75 76 // Cancel can be invoked to cancel the context before the test is completed. 77 // Tests which use the context to control goroutines and then wait for 78 // termination of those goroutines must call Cancel to avoid a deadlock. 79 // 80 // The cause, if non-empty, is turned into an error which is equivalend 81 // to context.Canceled. context.Cause will return that error for the 82 // context. 83 Cancel(cause string) 84 85 // Cleanup registers a callback that will get invoked when the test 86 // has finished. Callbacks get invoked in last-in-first-out order (LIFO). 87 // 88 // Beware of context cancellation. The following cleanup code 89 // will use a canceled context, which is not desirable: 90 // 91 // tCtx.Cleanup(func() { /* do something with tCtx */ }) 92 // tCtx.Cancel() 93 // 94 // A safer way to run cleanup code is: 95 // 96 // tCtx.CleanupCtx(func (tCtx ktesting.TContext) { /* do something with cleanup tCtx */ }) 97 Cleanup(func()) 98 99 // CleanupCtx is an alternative for Cleanup. The callback is passed a 100 // new TContext with the same logger and clients as the one CleanupCtx 101 // was invoked for. 102 CleanupCtx(func(TContext)) 103 104 // Expect wraps [gomega.Expect] such that a failure will be reported via 105 // [TContext.Fatal]. As with [gomega.Expect], additional values 106 // may get passed. Those values then all must be nil for the assertion 107 // to pass. This can be used with functions which return a value 108 // plus error: 109 // 110 // myAmazingThing := func(int, error) { ...} 111 // tCtx.Expect(myAmazingThing()).Should(gomega.Equal(1)) 112 Expect(actual interface{}, extra ...interface{}) gomega.Assertion 113 114 // ExpectNoError asserts that no error has occurred. 115 // 116 // As in [gomega], the optional explanation can be: 117 // - a [fmt.Sprintf] format string plus its argument 118 // - a function returning a string, which will be called 119 // lazy to construct the explanation if needed 120 // 121 // If an explanation is provided, then it replaces the default "Unexpected 122 // error" in the failure message. It's combined with additional details by 123 // adding a colon at the end, as when wrapping an error. Therefore it should 124 // not end with a punctuation mark or line break. 125 // 126 // Using ExpectNoError instead of the corresponding Gomega or testify 127 // assertions has the advantage that the failure message is short (good for 128 // aggregation in https://go.k8s.io/triage) with more details captured in the 129 // test log output (good when investigating one particular failure). 130 ExpectNoError(err error, explain ...interface{}) 131 132 // Logger returns a logger for the current test. This is a shortcut 133 // for calling klog.FromContext. 134 // 135 // Output emitted via this logger and the TB interface (like Logf) 136 // is formatted consistently. The TB interface generates a single 137 // message string, while Logger enables structured logging and can 138 // be passed down into code which expects a logger. 139 // 140 // To skip intermediate helper functions during stack unwinding, 141 // TB.Helper can be called in those functions. 142 Logger() klog.Logger 143 144 // TB returns the underlying TB. This can be used to "break the glass" 145 // and cast back into a testing.T or TB. Calling TB is necessary 146 // because TContext wraps the underlying TB. 147 TB() TB 148 149 // RESTConfig returns a config for a rest client with the UserAgent set 150 // to include the current test name or nil if not available. Several 151 // typed clients using this config are available through [Client], 152 // [Dynamic], [APIExtensions]. 153 RESTConfig() *rest.Config 154 155 RESTMapper() *restmapper.DeferredDiscoveryRESTMapper 156 Client() clientset.Interface 157 Dynamic() dynamic.Interface 158 APIExtensions() apiextensions.Interface 159 160 // The following methods must be implemented by every implementation 161 // of TContext to ensure that the leaf TContext is used, not some 162 // embedded TContext: 163 // - CleanupCtx 164 // - Expect 165 // - ExpectNoError 166 // - Logger 167 // 168 // Usually these methods would be stand-alone functions with a TContext 169 // parameter. Offering them as methods simplifies the test code. 170 } 171 172 // TB is the interface common to [testing.T], [testing.B], [testing.F] and 173 // [github.com/onsi/ginkgo/v2]. In contrast to [testing.TB], it can be 174 // implemented also outside of the testing package. 175 type TB interface { 176 Cleanup(func()) 177 Error(args ...any) 178 Errorf(format string, args ...any) 179 Fail() 180 FailNow() 181 Failed() bool 182 Fatal(args ...any) 183 Fatalf(format string, args ...any) 184 Helper() 185 Log(args ...any) 186 Logf(format string, args ...any) 187 Name() string 188 Setenv(key, value string) 189 Skip(args ...any) 190 SkipNow() 191 Skipf(format string, args ...any) 192 Skipped() bool 193 TempDir() string 194 } 195 196 // ContextTB adds support for cleanup callbacks with explicit context 197 // parameter. This is used when integrating with Ginkgo: then CleanupCtx 198 // gets implemented via ginkgo.DeferCleanup. 199 type ContextTB interface { 200 TB 201 CleanupCtx(func(ctx context.Context)) 202 } 203 204 // Init can be called in a unit or integration test to create 205 // a test context which: 206 // - has a per-test logger with verbosity derived from the -v command line flag 207 // - gets canceled when the test finishes (via [TB.Cleanup]) 208 // 209 // Note that the test context supports the interfaces of [TB] and 210 // [context.Context] and thus can be used like one of those where needed. 211 // It also has additional methods for retrieving the logger and canceling 212 // the context early, which can be useful in tests which want to wait 213 // for goroutines to terminate after cancellation. 214 // 215 // If the [TB] implementation also implements [ContextTB], then 216 // [TContext.CleanupCtx] uses [ContextTB.CleanupCtx] and uses 217 // the context passed into that callback. This can be used to let 218 // Ginkgo create a fresh context for cleanup code. 219 // 220 // Can be called more than once per test to get different contexts with 221 // independent cancellation. The default behavior describe above can be 222 // modified via optional functional options defined in [initoption]. 223 func Init(tb TB, opts ...InitOption) TContext { 224 tb.Helper() 225 226 c := internal.InitConfig{ 227 PerTestOutput: true, 228 } 229 for _, opt := range opts { 230 opt(&c) 231 } 232 233 // We don't need a Deadline implementation, testing.B doesn't have it. 234 // But if we have one, we'll use it to set a timeout shortly before 235 // the deadline. This needs to come before we wrap tb. 236 deadlineTB, deadlineOK := tb.(interface { 237 Deadline() (time.Time, bool) 238 }) 239 240 ctx := interruptCtx 241 if c.PerTestOutput { 242 config := ktesting.NewConfig( 243 ktesting.AnyToString(func(v interface{}) string { 244 return format.Object(v, 1) 245 }), 246 ktesting.VerbosityFlagName("v"), 247 ktesting.VModuleFlagName("vmodule"), 248 ) 249 250 // Copy klog settings instead of making the ktesting logger 251 // configurable directly. 252 var fs flag.FlagSet 253 config.AddFlags(&fs) 254 for _, name := range []string{"v", "vmodule"} { 255 from := flag.CommandLine.Lookup(name) 256 to := fs.Lookup(name) 257 if err := to.Value.Set(from.Value.String()); err != nil { 258 panic(err) 259 } 260 } 261 262 // Ensure consistent logging: this klog.Logger writes to tb, adding the 263 // date/time header, and our own wrapper emulates that behavior for 264 // Log/Logf/... 265 logger := ktesting.NewLogger(tb, config) 266 ctx = klog.NewContext(interruptCtx, logger) 267 268 tb = withKlogHeader(tb) 269 } 270 271 if deadlineOK { 272 if deadline, ok := deadlineTB.Deadline(); ok { 273 timeLeft := time.Until(deadline) 274 timeLeft -= CleanupGracePeriod 275 ctx, cancel := withTimeout(ctx, tb, timeLeft, fmt.Sprintf("test suite deadline (%s) is close, need to clean up before the %s cleanup grace period", deadline.Truncate(time.Second), CleanupGracePeriod)) 276 tCtx := tContext{ 277 Context: ctx, 278 testingTB: testingTB{TB: tb}, 279 cancel: cancel, 280 } 281 return tCtx 282 } 283 } 284 return WithCancel(InitCtx(ctx, tb)) 285 } 286 287 type InitOption = initoption.InitOption 288 289 // InitCtx is a variant of [Init] which uses an already existing context and 290 // whatever logger and timeouts are stored there. 291 // Functional options are part of the API, but currently 292 // there are none which have an effect. 293 func InitCtx(ctx context.Context, tb TB, _ ...InitOption) TContext { 294 tCtx := tContext{ 295 Context: ctx, 296 testingTB: testingTB{TB: tb}, 297 } 298 return tCtx 299 } 300 301 // WithTB constructs a new TContext with a different TB instance. 302 // This can be used to set up some of the context, in particular 303 // clients, in the root test and then run sub-tests: 304 // 305 // func TestSomething(t *testing.T) { 306 // tCtx := ktesting.Init(t) 307 // ... 308 // tCtx = ktesting.WithRESTConfig(tCtx, config) 309 // 310 // t.Run("sub", func (t *testing.T) { 311 // tCtx := ktesting.WithTB(tCtx, t) 312 // ... 313 // }) 314 // 315 // WithTB sets up cancellation for the sub-test. 316 func WithTB(parentCtx TContext, tb TB) TContext { 317 tCtx := InitCtx(parentCtx, tb) 318 tCtx = WithCancel(tCtx) 319 tCtx = WithClients(tCtx, 320 parentCtx.RESTConfig(), 321 parentCtx.RESTMapper(), 322 parentCtx.Client(), 323 parentCtx.Dynamic(), 324 parentCtx.APIExtensions(), 325 ) 326 return tCtx 327 } 328 329 // WithContext constructs a new TContext with a different Context instance. 330 // This can be used in callbacks which receive a Context, for example 331 // from Gomega: 332 // 333 // gomega.Eventually(tCtx, func(ctx context.Context) { 334 // tCtx := ktesting.WithContext(tCtx, ctx) 335 // ... 336 // 337 // This is important because the Context in the callback could have 338 // a different deadline than in the parent TContext. 339 func WithContext(parentCtx TContext, ctx context.Context) TContext { 340 tCtx := InitCtx(ctx, parentCtx.TB()) 341 tCtx = WithClients(tCtx, 342 parentCtx.RESTConfig(), 343 parentCtx.RESTMapper(), 344 parentCtx.Client(), 345 parentCtx.Dynamic(), 346 parentCtx.APIExtensions(), 347 ) 348 return tCtx 349 } 350 351 // WithValue wraps context.WithValue such that the result is again a TContext. 352 func WithValue(parentCtx TContext, key, val any) TContext { 353 ctx := context.WithValue(parentCtx, key, val) 354 return WithContext(parentCtx, ctx) 355 } 356 357 type tContext struct { 358 context.Context 359 testingTB 360 cancel func(cause string) 361 } 362 363 // testingTB is needed to avoid a name conflict 364 // between field and method in tContext. 365 type testingTB struct { 366 TB 367 } 368 369 func (tCtx tContext) Cancel(cause string) { 370 if tCtx.cancel != nil { 371 tCtx.cancel(cause) 372 } 373 } 374 375 func (tCtx tContext) CleanupCtx(cb func(TContext)) { 376 tCtx.Helper() 377 cleanupCtx(tCtx, cb) 378 } 379 380 func (tCtx tContext) Expect(actual interface{}, extra ...interface{}) gomega.Assertion { 381 tCtx.Helper() 382 return expect(tCtx, actual, extra...) 383 } 384 385 func (tCtx tContext) ExpectNoError(err error, explain ...interface{}) { 386 tCtx.Helper() 387 expectNoError(tCtx, err, explain...) 388 } 389 390 func cleanupCtx(tCtx TContext, cb func(TContext)) { 391 tCtx.Helper() 392 393 if tb, ok := tCtx.TB().(ContextTB); ok { 394 // Use context from base TB (most likely Ginkgo). 395 tb.CleanupCtx(func(ctx context.Context) { 396 tCtx := WithContext(tCtx, ctx) 397 cb(tCtx) 398 }) 399 return 400 } 401 402 tCtx.Cleanup(func() { 403 // Use new context. This is the code path for "go test". The 404 // context then has *no* deadline. In the code path above for 405 // Ginkgo, Ginkgo is more sophisticated and also applies 406 // timeouts to cleanup calls which accept a context. 407 childCtx := WithContext(tCtx, context.WithoutCancel(tCtx)) 408 cb(childCtx) 409 }) 410 } 411 412 func (tCtx tContext) Logger() klog.Logger { 413 return klog.FromContext(tCtx) 414 } 415 416 func (tCtx tContext) Error(args ...any) { 417 tCtx.Helper() 418 args = append([]any{"ERROR:"}, args...) 419 tCtx.testingTB.Error(args...) 420 } 421 422 func (tCtx tContext) Errorf(format string, args ...any) { 423 tCtx.Helper() 424 error := fmt.Sprintf(format, args...) 425 error = "ERROR: " + error 426 tCtx.testingTB.Error(error) 427 } 428 429 func (tCtx tContext) Fatal(args ...any) { 430 tCtx.Helper() 431 args = append([]any{"FATAL ERROR:"}, args...) 432 tCtx.testingTB.Fatal(args...) 433 } 434 435 func (tCtx tContext) Fatalf(format string, args ...any) { 436 tCtx.Helper() 437 error := fmt.Sprintf(format, args...) 438 error = "FATAL ERROR: " + error 439 tCtx.testingTB.Fatal(error) 440 } 441 442 func (tCtx tContext) TB() TB { 443 // Might have to unwrap twice, depending on how 444 // this tContext was constructed. 445 tb := tCtx.testingTB.TB 446 if k, ok := tb.(klogTB); ok { 447 return k.TB 448 } 449 return tb 450 } 451 452 func (tCtx tContext) RESTConfig() *rest.Config { 453 return nil 454 } 455 456 func (tCtx tContext) RESTMapper() *restmapper.DeferredDiscoveryRESTMapper { 457 return nil 458 } 459 460 func (tCtx tContext) Client() clientset.Interface { 461 return nil 462 } 463 464 func (tCtx tContext) Dynamic() dynamic.Interface { 465 return nil 466 } 467 468 func (tCtx tContext) APIExtensions() apiextensions.Interface { 469 return nil 470 } 471