From 20d9229c1959112eec44693192e844323493ab7e Mon Sep 17 00:00:00 2001
From: Patrick Ohly <patrick.ohly@intel.com>
Date: Wed, 29 Jun 2022 13:23:07 +0200
Subject: [PATCH] ktesting: add Stop

When testing.T.Log gets called after the test has completed, it panics. There's
also a data race (https://github.com/kubernetes/kubernetes/issues/110854).

Normally that should never happen because tests should ensure that all
goroutines have stopped before returning. But sometimes it is not possible to
do that. For those cases, "defer Stop(logger)" may be added to a test.  When
called, it will cause all future usage of the testing.T instance to be skipped.
---
 ktesting/testinglogger.go      | 55 ++++++++++++++++++++++++++++++++++
 ktesting/testinglogger_test.go | 25 ++++++++++++++++
 2 files changed, 80 insertions(+)

diff --git a/ktesting/testinglogger.go b/ktesting/testinglogger.go
index 905c1d634..db2a5990a 100644
--- a/ktesting/testinglogger.go
+++ b/ktesting/testinglogger.go
@@ -43,6 +43,7 @@ package ktesting
 
 import (
 	"bytes"
+	"fmt"
 	"strings"
 	"sync"
 	"time"
@@ -80,12 +81,22 @@ var _TL = NopTL{}
 
 // NewLogger constructs a new logger for the given test interface.
 //
+// Beware that testing.T and therefore this logger do not support logging after
+// a test has completed: then testing.T.Log has a data race and panics.
+//
+// A test therefore must ensure that all goroutines that might use the logger
+// have completed before returning or disable logging before returning with:
+//
+//    logger, ctx := ktesting.NewTestContext(t)
+//    defer ktesting.Stop(logger)
+//
 // Experimental
 //
 // Notice: This type is EXPERIMENTAL and may be changed or removed in a
 // later release.
 func NewLogger(t TL, c *Config) logr.Logger {
 	return logr.New(&tlogger{
+		mutex:  new(sync.Mutex),
 		t:      t,
 		prefix: "",
 		values: nil,
@@ -94,6 +105,25 @@ func NewLogger(t TL, c *Config) logr.Logger {
 	})
 }
 
+// Stop ensures that all future log entries get discarded. This is necessary
+// when goroutines might continue to use it after the test is completed because
+// testing.T.Log panics when called after test completion.
+//
+// Stop may only be called for a logger returned by NewLogger or NewTestContext.
+// Calling it for other loggers is an error and triggers a panic.
+//
+// Experimental
+//
+// Notice: This function is EXPERIMENTAL and may be changed or removed in a
+// later release.
+func Stop(logger logr.Logger) {
+	l, ok := logger.GetSink().(*tlogger)
+	if !ok {
+		panic(fmt.Sprintf("Stop should only be called for a logger created by ktesting, got called for %T.", logger.GetSink()))
+	}
+	l.stop()
+}
+
 // Buffer stores log entries as formatted text and structured data.
 // It is safe to use this concurrently.
 //
@@ -228,6 +258,7 @@ func (b *buffer) Data() Log {
 }
 
 type tlogger struct {
+	mutex  *sync.Mutex
 	t      TL
 	prefix string
 	values []interface{}
@@ -239,10 +270,22 @@ func (l *tlogger) Init(info logr.RuntimeInfo) {
 }
 
 func (l *tlogger) GetCallStackHelper() func() {
+	l.mutex.Lock()
+	defer l.mutex.Unlock()
+	if l.t == nil {
+		return func() {}
+	}
+
 	return l.t.Helper
 }
 
 func (l *tlogger) Info(level int, msg string, kvList ...interface{}) {
+	l.mutex.Lock()
+	defer l.mutex.Unlock()
+	if l.t == nil {
+		return
+	}
+
 	l.t.Helper()
 	buffer := &bytes.Buffer{}
 	merged := serialize.MergeKVs(l.values, kvList)
@@ -255,6 +298,12 @@ func (l *tlogger) Enabled(level int) bool {
 }
 
 func (l *tlogger) Error(err error, msg string, kvList ...interface{}) {
+	l.mutex.Lock()
+	defer l.mutex.Unlock()
+	if l.t == nil {
+		return
+	}
+
 	l.t.Helper()
 	buffer := &bytes.Buffer{}
 	if err != nil {
@@ -307,6 +356,12 @@ func (l *tlogger) log(what LogType, msg string, level int, buffer *bytes.Buffer,
 	)
 }
 
+func (l *tlogger) stop() {
+	l.mutex.Lock()
+	defer l.mutex.Unlock()
+	l.t = nil
+}
+
 // WithName returns a new logr.Logger with the specified name appended.  klogr
 // uses '/' characters to separate name elements.  Callers should not pass '/'
 // in the provided name string, but this library does not actually enforce that.
diff --git a/ktesting/testinglogger_test.go b/ktesting/testinglogger_test.go
index c4efa4629..9dbc8363a 100644
--- a/ktesting/testinglogger_test.go
+++ b/ktesting/testinglogger_test.go
@@ -11,6 +11,7 @@ import (
 	"bytes"
 	"errors"
 	"fmt"
+	"sync"
 	"testing"
 
 	"k8s.io/klog/v2/ktesting"
@@ -147,3 +148,27 @@ func (l *logToBuf) Log(args ...interface{}) {
 	}
 	l.Write([]byte("\n"))
 }
+
+func TestStop(t *testing.T) {
+	// This test is set up so that a subtest spawns a goroutine and that
+	// goroutine logs through ktesting *after* the subtest has
+	// completed. This is not supported by testing.T.Log and normally
+	// leads to:
+	//   panic: Log in goroutine after TestGoroutines/Sub has completed: INFO hello world
+	//
+	// It works with ktesting if (and only if) logging gets stopped
+	// before returning from the test.
+	var wg sync.WaitGroup
+	wg.Add(1)
+	t.Run("Sub", func(t *testing.T) {
+		logger, _ := ktesting.NewTestContext(t)
+		defer ktesting.Stop(logger)
+		go func() {
+			// Wait for test to have returned.
+			wg.Wait()
+			logger.Info("hello world")
+		}()
+	})
+	// Allow goroutine above to proceed.
+	wg.Done()
+}