testr: prevent panic when test keeps goroutines running

testr/testr.go

@@ -18,6 +18,7 @@ limitations under the License.
 package testr
 
 import (
+	"sync/atomic"
 	"testing"
 
 	"github.com/go-logr/logr"
@@ -57,6 +58,12 @@ type TestingT interface {
 	Log(args ...interface{})
 }
 
+// cleanupT is implemented by the t instances that we get, but we don't require
+// it in our API.
+type cleanupT interface {
+	Cleanup(func())
+}
+
 // NewWithInterface returns a logr.Logger that prints through a
 // TestingT object.
 // In contrast to the simpler New, output formatting can be configured.
@@ -65,20 +72,32 @@ func NewWithInterface(t TestingT, opts Options) logr.Logger {
 	return logr.New(&l)
 }
 
+var noT TestingT
+
 func newLoggerInterfaceWithOptions(t TestingT, opts Options) testloggerInterface {
-	return testloggerInterface{
-		t: t,
+	l := testloggerInterface{
+		t: new(atomic.Value),
 		Formatter: funcr.NewFormatter(funcr.Options{
 			LogTimestamp: opts.LogTimestamp,
 			Verbosity:    opts.Verbosity,
 		}),
 	}
+	l.t.Store(&t)
+	if cleanup, ok := t.(cleanupT); ok {
+		cleanup.Cleanup(func() {
+			// Cannot store nil.
+			l.t.Store(&noT)
+		})
+	}
+	return l
 }
 
 // Underlier exposes access to the underlying testing.T instance. Since
 // callers only have a logr.Logger, they have to know which
 // implementation is in use, so this interface is less of an
 // abstraction and more of a way to test type conversion.
+//
+// Returns nil when the test has already completed.
 type Underlier interface {
 	GetUnderlying() *testing.T
 }
@@ -87,6 +106,8 @@ type Underlier interface {
 // callers only have a logr.Logger, they have to know which
 // implementation is in use, so this interface is less of an
 // abstraction and more of a way to test type conversion.
+//
+// Returns nil when the test has already completed.
 type UnderlierInterface interface {
 	GetUnderlying() TestingT
 }
@@ -118,15 +139,24 @@ type testlogger struct {
 }
 
 func (l testlogger) GetUnderlying() *testing.T {
+	t := l.testloggerInterface.GetUnderlying()
+	if t == nil {
+		return nil
+	}
+
 	// This method is defined on testlogger, so the only type this could
 	// possibly be is testing.T, even though that's not guaranteed by the type
 	// system itself.
-	return l.t.(*testing.T) //nolint:forcetypeassert
+	return t.(*testing.T) //nolint:forcetypeassert
 }
 
 type testloggerInterface struct {
 	funcr.Formatter
-	t TestingT
+
+	// t will be cleared when the test is done because then t becomes
+	// unusable. In particular t.Log will panic with "Log in goroutine
+	// after ... has completed".
+	t *atomic.Value
 }
 
 func (l testloggerInterface) WithName(name string) logr.LogSink {
@@ -140,21 +170,41 @@ func (l testloggerInterface) WithValues(kvList ...interface{}) logr.LogSink {
 }
 
 func (l testloggerInterface) GetCallStackHelper() func() {
-	return l.t.Helper
+	t := l.GetUnderlying()
+	if t == nil {
+		return func() {}
+	}
+
+	return t.Helper
 }
 
 func (l testloggerInterface) Info(level int, msg string, kvList ...interface{}) {
-	l.t.Helper()
-	logInfo(l.t, l.FormatInfo, level, msg, kvList...)
+	t := l.GetUnderlying()
+	if t == nil {
+		return
+	}
+
+	t.Helper()
+	logInfo(t, l.FormatInfo, level, msg, kvList...)
 }
 
 func (l testloggerInterface) Error(err error, msg string, kvList ...interface{}) {
-	l.t.Helper()
-	logError(l.t, l.FormatError, err, msg, kvList...)
+	t := l.GetUnderlying()
+	if t == nil {
+		return
+	}
+
+	t.Helper()
+	logError(t, l.FormatError, err, msg, kvList...)
 }
 
 func (l testloggerInterface) GetUnderlying() TestingT {
-	return l.t
+	t := l.t.Load()
+	if t == &noT {
+		return nil
+	}
+	// This is the only type that this could possibly be.
+	return *t.(*TestingT) //nolint:forcetypeassert
 }
 
 // Assert conformance to the interfaces.

testr/testr_test.go

@@ -18,6 +18,8 @@ package testr
 
 import (
 	"fmt"
+	"os"
+	"sync"
 	"testing"
 
 	"github.com/go-logr/logr"
@@ -87,3 +89,42 @@ func helper2(log logr.Logger, msg string) {
 	helper()
 	log.Info(msg)
 }
+
+var testStopWG sync.WaitGroup
+
+func TestMain(m *testing.M) {
+	exitCode := m.Run()
+	testStopWG.Wait()
+
+	os.Exit(exitCode)
+}
+
+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 testr if (and only if) logging gets discarded
+	// before returning from the test.
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	testStopWG.Add(1)
+	logger := New(t)
+	go func() {
+		defer testStopWG.Done()
+
+		// Wait for test to have returned.
+		wg.Wait()
+
+		// This output must be discarded because the test has
+		// completed.
+		logger.Info("simple info message")
+		logger.Error(nil, "error message")
+		logger.WithName("me").WithValues("completed", true).Info("complex info message", "anotherValue", 1)
+	}()
+	// Allow goroutine above to proceed.
+	wg.Done()
+}