From c665bbabb360e855c60a2e53f05087bf6aa724d5 Mon Sep 17 00:00:00 2001 From: Prashant Varanasi Date: Tue, 1 Sep 2020 13:29:58 -0700 Subject: [PATCH] Add options to customize Fatal behaviour for better testability (#861) Fixes #846. There's currently no easy way to test Fatal from outside of zap, as it triggers an os.Exit(1). Add options to customize the behaviour (allowing for panic, or Goexit), which can be used with `recover()` in tests. --- logger.go | 9 ++++++- logger_test.go | 45 +++++++++++++++++++++++++++++++++++ options.go | 7 ++++++ zapcore/entry.go | 5 ++++ zapcore/entry_test.go | 55 +++++++++++++++++++++++++++++++------------ 5 files changed, 105 insertions(+), 16 deletions(-) diff --git a/logger.go b/logger.go index 4726456d7..ea484aed1 100644 --- a/logger.go +++ b/logger.go @@ -49,6 +49,7 @@ type Logger struct { addStack zapcore.LevelEnabler callerSkip int + onFatal zapcore.CheckWriteAction // default is WriteThenFatal } // New constructs a new Logger from the provided zapcore.Core and Options. If @@ -280,7 +281,13 @@ func (log *Logger) check(lvl zapcore.Level, msg string) *zapcore.CheckedEntry { case zapcore.PanicLevel: ce = ce.Should(ent, zapcore.WriteThenPanic) case zapcore.FatalLevel: - ce = ce.Should(ent, zapcore.WriteThenFatal) + onFatal := log.onFatal + // Noop is the default value for CheckWriteAction, and it leads to + // continued execution after a Fatal which is unexpected. + if onFatal == zapcore.WriteThenNoop { + onFatal = zapcore.WriteThenFatal + } + ce = ce.Should(ent, onFatal) case zapcore.DPanicLevel: if log.development { ce = ce.Should(ent, zapcore.WriteThenPanic) diff --git a/logger_test.go b/logger_test.go index 6992a9a19..0fb0600e7 100644 --- a/logger_test.go +++ b/logger_test.go @@ -535,6 +535,51 @@ func TestLoggerConcurrent(t *testing.T) { }) } +func TestLoggerCustomOnFatal(t *testing.T) { + tests := []struct { + msg string + onFatal zapcore.CheckWriteAction + recoverValue interface{} + }{ + { + msg: "panic", + onFatal: zapcore.WriteThenPanic, + recoverValue: "fatal", + }, + { + msg: "goexit", + onFatal: zapcore.WriteThenGoexit, + recoverValue: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.msg, func(t *testing.T) { + withLogger(t, InfoLevel, opts(OnFatal(tt.onFatal)), func(logger *Logger, logs *observer.ObservedLogs) { + + var finished bool + recovered := make(chan interface{}) + go func() { + defer func() { + recovered <- recover() + }() + + logger.Fatal("fatal") + finished = true + }() + + assert.Equal(t, tt.recoverValue, <-recovered, "unexpected value from recover()") + assert.False(t, finished, "expect goroutine to not finish after Fatal") + + assert.Equal(t, []observer.LoggedEntry{{ + Entry: zapcore.Entry{Level: FatalLevel, Message: "fatal"}, + Context: []Field{}, + }}, logs.AllUntimed(), "unexpected logs") + }) + }) + } +} + func infoLog(logger *Logger, msg string, fields ...Field) { logger.Info(msg, fields...) } diff --git a/options.go b/options.go index 2bf2cb222..0135c2092 100644 --- a/options.go +++ b/options.go @@ -131,3 +131,10 @@ func IncreaseLevel(lvl zapcore.LevelEnabler) Option { } }) } + +// OnFatal sets the action to take on fatal logs. +func OnFatal(action zapcore.CheckWriteAction) Option { + return optionFunc(func(log *Logger) { + log.onFatal = action + }) +} diff --git a/zapcore/entry.go b/zapcore/entry.go index 717dc0c4f..4aa8b4f90 100644 --- a/zapcore/entry.go +++ b/zapcore/entry.go @@ -22,6 +22,7 @@ package zapcore import ( "fmt" + "runtime" "strings" "sync" "time" @@ -159,6 +160,8 @@ const ( // WriteThenNoop indicates that nothing special needs to be done. It's the // default behavior. WriteThenNoop CheckWriteAction = iota + // WriteThenGoexit runs runtime.Goexit after Write. + WriteThenGoexit // WriteThenPanic causes a panic after Write. WriteThenPanic // WriteThenFatal causes a fatal os.Exit after Write. @@ -231,6 +234,8 @@ func (ce *CheckedEntry) Write(fields ...Field) { panic(msg) case WriteThenFatal: exit.Exit() + case WriteThenGoexit: + runtime.Goexit() } } diff --git a/zapcore/entry_test.go b/zapcore/entry_test.go index 569c4e1e0..4c2d67eae 100644 --- a/zapcore/entry_test.go +++ b/zapcore/entry_test.go @@ -29,6 +29,22 @@ import ( "github.com/stretchr/testify/assert" ) +func assertGoexit(t *testing.T, f func()) { + var finished bool + recovered := make(chan interface{}) + go func() { + defer func() { + recovered <- recover() + }() + + f() + finished = true + }() + + assert.Nil(t, <-recovered, "Goexit should cause recover to return nil") + assert.False(t, finished, "Goroutine should not finish after Goexit") +} + func TestPutNilEntry(t *testing.T) { // Pooling nil entries defeats the purpose. var wg sync.WaitGroup @@ -88,20 +104,29 @@ func TestEntryCaller(t *testing.T) { } func TestCheckedEntryWrite(t *testing.T) { - // Nil checked entries are safe. - var ce *CheckedEntry - assert.NotPanics(t, func() { ce.Write() }, "Unexpected panic writing nil CheckedEntry.") - - // WriteThenPanic - ce = ce.Should(Entry{}, WriteThenPanic) - assert.Panics(t, func() { ce.Write() }, "Expected to panic when WriteThenPanic is set.") - ce.reset() - - // WriteThenFatal - ce = ce.Should(Entry{}, WriteThenFatal) - stub := exit.WithStub(func() { - ce.Write() + t.Run("nil is safe", func(t *testing.T) { + var ce *CheckedEntry + assert.NotPanics(t, func() { ce.Write() }, "Unexpected panic writing nil CheckedEntry.") + }) + + t.Run("WriteThenPanic", func(t *testing.T) { + var ce *CheckedEntry + ce = ce.Should(Entry{}, WriteThenPanic) + assert.Panics(t, func() { ce.Write() }, "Expected to panic when WriteThenPanic is set.") + }) + + t.Run("WriteThenGoexit", func(t *testing.T) { + var ce *CheckedEntry + ce = ce.Should(Entry{}, WriteThenGoexit) + assertGoexit(t, func() { ce.Write() }) + }) + + t.Run("WriteThenFatal", func(t *testing.T) { + var ce *CheckedEntry + ce = ce.Should(Entry{}, WriteThenFatal) + stub := exit.WithStub(func() { + ce.Write() + }) + assert.True(t, stub.Exited, "Expected to exit when WriteThenFatal is set.") }) - assert.True(t, stub.Exited, "Expected to exit when WriteThenFatal is set.") - ce.reset() }