Add MuteAfterTestCompletion() option for zaptest

zaptest/logger.go

@@ -22,6 +22,7 @@ package zaptest
 
 import (
 	"bytes"
+	"strings"
 
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
@@ -33,8 +34,9 @@ type LoggerOption interface {
 }
 
 type loggerOptions struct {
-	Level      zapcore.LevelEnabler
-	zapOptions []zap.Option
+	Level                   zapcore.LevelEnabler
+	MuteAfterTestCompletion bool
+	zapOptions              []zap.Option
 }
 
 type loggerOptionFunc func(*loggerOptions)
@@ -58,6 +60,12 @@ func WrapOptions(zapOpts ...zap.Option) LoggerOption {
 	})
 }
 
+func MuteAfterTestCompletion() LoggerOption {
+	return loggerOptionFunc(func(opts *loggerOptions) {
+		opts.MuteAfterTestCompletion = true
+	})
+}
+
 // NewLogger builds a new Logger that logs all messages to the given
 // testing.TB.
 //
@@ -83,6 +91,11 @@ func NewLogger(t TestingT, opts ...LoggerOption) *zap.Logger {
 	}
 
 	writer := NewTestingWriter(t)
+
+	if cfg.MuteAfterTestCompletion {
+		writer.muteAfterTestCompletion = true
+	}
+
 	zapOptions := []zap.Option{
 		// Send zap errors to the same writer and mark the test as failed if
 		// that happens.
@@ -107,6 +120,13 @@ type TestingWriter struct {
 	// If true, the test will be marked as failed if this TestingWriter is
 	// ever used.
 	markFailed bool
+
+	// If true, we want to mute logging after the test has completed.
+	// We detect this by catching the panic
+	muteAfterTestCompletion bool
+
+	// If true, we've muted due to a prior panic
+	muted bool
 }
 
 // NewTestingWriter builds a new TestingWriter that writes to the given
@@ -139,6 +159,26 @@ func (w TestingWriter) WithMarkFailed(v bool) TestingWriter {
 func (w TestingWriter) Write(p []byte) (n int, err error) {
 	n = len(p)
 
+	if w.muted {
+		// Early exit if we've already muted
+		return n, nil
+	}
+
+	if w.muteAfterTestCompletion {
+		// Set up to catch the panic that happens if the test has completed
+		defer func() {
+			if r := recover(); r != nil {
+				if s, ok := r.(string); ok && strings.HasPrefix(s, "Log in goroutine after") {
+					w.muted = true
+					return
+				}
+
+				// Re-panic if it's not the expected panic (just in case)
+				panic(r)
+			}
+		}()
+	}
+
 	// Strip trailing newline because t.Log always adds one.
 	p = bytes.TrimRight(p, "\n")
 

zaptest/logger_test.go

@@ -191,3 +191,56 @@ func (t *testLogSpy) AssertFailed() {
 func (t *testLogSpy) assertFailed(v bool, msg string) {
 	assert.Equal(t.TB, v, t.failed, msg)
 }
+
+func TestTestLoggerAfterTestCompletedPanics(t *testing.T) {
+	// A wrapper that always panics as though the wrapped test has completed.
+	w := newTestFinishedWrapper(t)
+
+	log := NewLogger(w)
+	assert.Panics(t, func() {
+		log.Info("foo")
+	})
+}
+
+func TestTestLoggerWithMutingAfterTestCompletedReturns(t *testing.T) {
+	// A wrapper that always panics as though the wrapped test has completed.
+	w := newTestFinishedWrapper(t)
+
+	log := NewLogger(w, MuteAfterTestCompletion())
+	log.Info("foo")
+}
+
+// testFinishedWrapper is a TestingT wrapper that panics if you try to log
+type testFinishedWrapper struct {
+	t *testing.T
+}
+
+func newTestFinishedWrapper(t *testing.T) *testFinishedWrapper {
+	return &testFinishedWrapper{
+		t: t,
+	}
+}
+
+func (f *testFinishedWrapper) Logf(string, ...interface{}) {
+	panic("Log in goroutine after " + f.t.Name() + " has completed")
+}
+
+func (f *testFinishedWrapper) Errorf(string, ...interface{}) {
+	panic("Error in goroutine after " + f.t.Name() + " has completed")
+}
+
+func (f *testFinishedWrapper) Fail() {
+	f.t.Fail()
+}
+
+func (f *testFinishedWrapper) Failed() bool {
+	return f.t.Failed()
+}
+
+func (f *testFinishedWrapper) Name() string {
+	return f.t.Name()
+}
+
+func (f *testFinishedWrapper) FailNow() {
+	f.t.FailNow()
+}