diff --git a/zaptest/logger.go b/zaptest/logger.go index 4734c33f6..8a448b1ce 100644 --- a/zaptest/logger.go +++ b/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") diff --git a/zaptest/logger_test.go b/zaptest/logger_test.go index 40e368b50..90f085906 100644 --- a/zaptest/logger_test.go +++ b/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() +}