Add options to customize Fatal behaviour for better testability

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,11 @@ 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
+		if onFatal == 0 {
+			onFatal = zapcore.WriteThenFatal
+		}
+		ce = ce.Should(ent, onFatal)
 	case zapcore.DPanicLevel:
 		if log.development {
 			ce = ce.Should(ent, zapcore.WriteThenPanic)

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...)
 }

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
+	})
+}

zapcore/entry.go

@@ -22,6 +22,7 @@ package zapcore
 
 import (
 	"fmt"
+	"runtime"
 	"strings"
 	"sync"
 	"time"
@@ -163,6 +164,8 @@ const (
 	WriteThenPanic
 	// WriteThenFatal causes a fatal os.Exit after Write.
 	WriteThenFatal
+	// WriteThenGoexit runs runtime.Goexit after Write.
+	WriteThenGoexit
 )
 
 // CheckedEntry is an Entry together with a collection of Cores that have
@@ -231,6 +234,8 @@ func (ce *CheckedEntry) Write(fields ...Field) {
 		panic(msg)
 	case WriteThenFatal:
 		exit.Exit()
+	case WriteThenGoexit:
+		runtime.Goexit()
 	}
 }
 

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()
 }