From 2be620216affd84620a08ee082d6e074d1bb5ca4 Mon Sep 17 00:00:00 2001
From: Albert Salim <albertlimca@gmail.com>
Date: Sat, 6 Oct 2018 18:08:19 +0800
Subject: [PATCH] Add option to panic in `test.NewNullLogger` to allow testing
 of calls to `Fatal*`

See #813
---
 alt_exit.go             |  2 ++
 entry.go                |  6 +++---
 hooks/test/test.go      | 14 +++++++++++++-
 hooks/test/test_test.go | 11 +++++++++++
 logger.go               |  9 ++++++---
 5 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/alt_exit.go b/alt_exit.go
index 8af9063..f1bb44c 100644
--- a/alt_exit.go
+++ b/alt_exit.go
@@ -45,6 +45,8 @@ func runHandlers() {
 	}
 }
 
+type exitFunc func(int)
+
 // Exit runs all the Logrus atexit handlers and then terminates the program using os.Exit(code)
 func Exit(code int) {
 	runHandlers()
diff --git a/entry.go b/entry.go
index ca634a6..225bdb6 100644
--- a/entry.go
+++ b/entry.go
@@ -198,7 +198,7 @@ func (entry *Entry) Fatal(args ...interface{}) {
 	if entry.Logger.IsLevelEnabled(FatalLevel) {
 		entry.log(FatalLevel, fmt.Sprint(args...))
 	}
-	Exit(1)
+	entry.Logger.Exit(1)
 }
 
 func (entry *Entry) Panic(args ...interface{}) {
@@ -246,7 +246,7 @@ func (entry *Entry) Fatalf(format string, args ...interface{}) {
 	if entry.Logger.IsLevelEnabled(FatalLevel) {
 		entry.Fatal(fmt.Sprintf(format, args...))
 	}
-	Exit(1)
+	entry.Logger.Exit(1)
 }
 
 func (entry *Entry) Panicf(format string, args ...interface{}) {
@@ -293,7 +293,7 @@ func (entry *Entry) Fatalln(args ...interface{}) {
 	if entry.Logger.IsLevelEnabled(FatalLevel) {
 		entry.Fatal(entry.sprintlnn(args...))
 	}
-	Exit(1)
+	entry.Logger.Exit(1)
 }
 
 func (entry *Entry) Panicln(args ...interface{}) {
diff --git a/hooks/test/test.go b/hooks/test/test.go
index 234a17d..f84fe80 100644
--- a/hooks/test/test.go
+++ b/hooks/test/test.go
@@ -39,12 +39,24 @@ func NewLocal(logger *logrus.Logger) *Hook {
 
 }
 
+type TestOption func(logger *logrus.Logger)
+
+func FatalPanics(logger *logrus.Logger) {
+	logger.Exit = func(code int) {
+		panic(code)
+	}
+}
+
 // NewNullLogger creates a discarding logger and installs the test hook.
-func NewNullLogger() (*logrus.Logger, *Hook) {
+func NewNullLogger(options ...TestOption) (*logrus.Logger, *Hook) {
 
 	logger := logrus.New()
 	logger.Out = ioutil.Discard
 
+	for _, option := range options {
+		option(logger)
+	}
+
 	return logger, NewLocal(logger)
 
 }
diff --git a/hooks/test/test_test.go b/hooks/test/test_test.go
index d6f6d30..692d36a 100644
--- a/hooks/test/test_test.go
+++ b/hooks/test/test_test.go
@@ -71,3 +71,14 @@ func TestLoggingWithHooksRace(t *testing.T) {
 	entries := hook.AllEntries()
 	assert.Equal(100, len(entries))
 }
+
+func TestFatalWithPanic(t *testing.T) {
+	assert := assert.New(t)
+
+	logger, hook := NewNullLogger(FatalPanics)
+
+	assert.Nil(hook.LastEntry())
+	assert.Equal(0, len(hook.Entries))
+
+	assert.Panics(func() { logger.Fatal("something went wrong") })
+}
diff --git a/logger.go b/logger.go
index b67bfcb..188c600 100644
--- a/logger.go
+++ b/logger.go
@@ -32,6 +32,8 @@ type Logger struct {
 	mu MutexWrap
 	// Reusable empty entry
 	entryPool sync.Pool
+	// Function to exit the application, defaults to `Exit()`
+	Exit exitFunc
 }
 
 type MutexWrap struct {
@@ -73,6 +75,7 @@ func New() *Logger {
 		Formatter: new(TextFormatter),
 		Hooks:     make(LevelHooks),
 		Level:     InfoLevel,
+		Exit:      Exit,
 	}
 }
 
@@ -173,7 +176,7 @@ func (logger *Logger) Fatalf(format string, args ...interface{}) {
 		entry.Fatalf(format, args...)
 		logger.releaseEntry(entry)
 	}
-	Exit(1)
+	logger.Exit(1)
 }
 
 func (logger *Logger) Panicf(format string, args ...interface{}) {
@@ -236,7 +239,7 @@ func (logger *Logger) Fatal(args ...interface{}) {
 		entry.Fatal(args...)
 		logger.releaseEntry(entry)
 	}
-	Exit(1)
+	logger.Exit(1)
 }
 
 func (logger *Logger) Panic(args ...interface{}) {
@@ -299,7 +302,7 @@ func (logger *Logger) Fatalln(args ...interface{}) {
 		entry.Fatalln(args...)
 		logger.releaseEntry(entry)
 	}
-	Exit(1)
+	logger.Exit(1)
 }
 
 func (logger *Logger) Panicln(args ...interface{}) {