From c2fbbd6947e4931ec5e8c540bbba4a0798674765 Mon Sep 17 00:00:00 2001
From: Pieter Claerhout <pieter@yellowduck.be>
Date: Thu, 6 Feb 2020 15:14:10 +0100
Subject: [PATCH] Fixed the testing and improved the CheckError function

---
 logger.go          | 24 ++++++++++++++++--------
 logger_internal.go |  5 ++---
 logger_test.go     |  4 ++--
 3 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/logger.go b/logger.go
index a4adbce..356ecc1 100644
--- a/logger.go
+++ b/logger.go
@@ -2,6 +2,7 @@ package log
 
 import (
 	"fmt"
+	"io"
 	"os"
 	"strings"
 	"time"
@@ -31,10 +32,10 @@ var DebugSQLMode = false
 var TimeZone *time.Location
 
 // Stdout is the writer to where the stdout messages should be written (defaults to os.Stdout)
-var Stdout *os.File = os.Stdout
+var Stdout io.Writer = os.Stdout
 
 // Stderr is the writer to where the stderr messages should be written (defaults to os.Stderr)
-var Stderr *os.File = os.Stderr
+var Stderr io.Writer = os.Stderr
 
 // DefaultTimeFormat is the default format to use for the timestamps
 var DefaultTimeFormat = "2006-01-02 15:04:05.000"
@@ -222,11 +223,18 @@ func Fatalf(format string, args ...interface{}) {
 //
 // If DebugMode is enabled a stack trace will also be printed to stderr
 func CheckError(err error) {
-	if err != nil {
-		printMessage("FATAL", err.Error())
-		if DebugMode {
-			StackTrace(err)
-		}
-		OsExit(1)
+
+	if err == nil {
+		return
+	}
+
+	msg := err.Error()
+	if DebugMode {
+		msg = formatMessage(FormattedStackTrace(err))
 	}
+
+	printMessage("FATAL", msg)
+
+	OsExit(1)
+
 }
diff --git a/logger_internal.go b/logger_internal.go
index 8ec07dd..0b7e96f 100644
--- a/logger_internal.go
+++ b/logger_internal.go
@@ -8,7 +8,6 @@ import (
 	"time"
 
 	"github.com/fatih/color"
-	"github.com/mattn/go-colorable"
 )
 
 var logMutex = &sync.Mutex{}
@@ -66,9 +65,9 @@ func printMessage(level string, message string) {
 	}
 
 	if PrintColors {
-		cw := colorable.NewColorable(w)
+		// cw := colorable.NewColorable(w)
 		if c, ok := colors[level]; ok {
-			c.Fprintln(cw, message)
+			c.Fprintln(w, message)
 			return
 		}
 	}
diff --git a/logger_test.go b/logger_test.go
index 2804cc2..653af0a 100644
--- a/logger_test.go
+++ b/logger_test.go
@@ -564,7 +564,7 @@ func TestCheckError(t *testing.T) {
 		{"nil-debug-color", nil, true, "", "", -1},
 
 		{"err-release-nocolor", errors.New("test"), false, "", "test | FATAL | test\n", 1},
-		{"err-debug-nocolor", errors.New("test"), true, "", "test | FATAL | test\ntest | ERROR | *errors.fundamental test\n", 1},
+		{"err-debug-nocolor", errors.New("test"), true, "", "test | FATAL | *errors.fundamental test\n", 1},
 	}
 
 	for _, tc := range tests {
@@ -593,7 +593,7 @@ func TestCheckError(t *testing.T) {
 
 			assert.Equal(t, tc.expectedStdout, actualStdOut)
 			if tc.debug {
-				assert.True(t, strings.HasPrefix(actualStdErr, tc.expectedStderr))
+				assert.True(t, strings.HasPrefix(actualStdErr, tc.expectedStderr), actualStdErr)
 			} else {
 				assert.Equal(t, tc.expectedStderr, actualStdErr)
 			}