Skip to content

Commit

Permalink
Add options to customize Fatal behaviour for better testability (#861)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
prashantv authored Sep 1, 2020
1 parent be2be86 commit 217b2cb
Show file tree
Hide file tree
Showing 5 changed files with 105 additions and 16 deletions.
9 changes: 8 additions & 1 deletion logger.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down
45 changes: 45 additions & 0 deletions logger_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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...)
}
Expand Down
7 changes: 7 additions & 0 deletions options.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
})
}
5 changes: 5 additions & 0 deletions zapcore/entry.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ package zapcore

import (
"fmt"
"runtime"
"strings"
"sync"
"time"
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -231,6 +234,8 @@ func (ce *CheckedEntry) Write(fields ...Field) {
panic(msg)
case WriteThenFatal:
exit.Exit()
case WriteThenGoexit:
runtime.Goexit()
}
}

Expand Down
55 changes: 40 additions & 15 deletions zapcore/entry_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()
}

0 comments on commit 217b2cb

Please sign in to comment.