From 8982f5d70236f6083740c3de66ae2a58607eb260 Mon Sep 17 00:00:00 2001
From: IamCathal <cathalsocallaghan@gmail.com>
Date: Wed, 9 Sep 2020 00:23:57 +0100
Subject: [PATCH] Add unit tests for custom log formatter

Also moved is Colorful check to WriteMsg function to make the interface for user's using the custom logging formatting simpler. The user does not have to check if the text is colorful now, the WriteMsg function handles it.
---
 pkg/logs/console.go                           | 10 ++----
 .../logformattertest/log_formatter_test.go    | 36 +++++++++++++++++++
 2 files changed, 39 insertions(+), 7 deletions(-)
 create mode 100644 pkg/logs/logformattertest/log_formatter_test.go

diff --git a/pkg/logs/console.go b/pkg/logs/console.go
index 34114e4a..a3e5fb5a 100644
--- a/pkg/logs/console.go
+++ b/pkg/logs/console.go
@@ -58,10 +58,6 @@ type consoleWriter struct {
 func (c *consoleWriter) Format(lm *LogMsg) string {
 	msg := lm.Msg
 
-	if c.Colorful {
-		msg = strings.Replace(lm.Msg, levelPrefix[lm.Level], colors[lm.Level](levelPrefix[lm.Level]), 1)
-	}
-
 	h, _, _ := formatTimeHeader(lm.When)
 	bytes := append(append(h, msg...), '\n')
 
@@ -105,13 +101,13 @@ func (c *consoleWriter) WriteMsg(lm *LogMsg) error {
 	if lm.Level > c.Level {
 		return nil
 	}
-	// fmt.Printf("Formatted: %s\n\n", c.fmtter.Format(lm))
+
+	msg := ""
+
 	if c.Colorful {
 		lm.Msg = strings.Replace(lm.Msg, levelPrefix[lm.Level], colors[lm.Level](levelPrefix[lm.Level]), 1)
 	}
 
-	msg := ""
-
 	if c.customFormatter != nil {
 		msg = c.customFormatter(lm)
 	} else {
diff --git a/pkg/logs/logformattertest/log_formatter_test.go b/pkg/logs/logformattertest/log_formatter_test.go
new file mode 100644
index 00000000..2d99a8e6
--- /dev/null
+++ b/pkg/logs/logformattertest/log_formatter_test.go
@@ -0,0 +1,36 @@
+package logformattertest
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/astaxie/beego/pkg/common"
+	"github.com/astaxie/beego/pkg/logs"
+)
+
+func customFormatter(lm *logs.LogMsg) string {
+	return fmt.Sprintf("[CUSTOM CONSOLE LOGGING] %s", lm.Msg)
+}
+
+func globalFormatter(lm *logs.LogMsg) string {
+	return fmt.Sprintf("[GLOBAL] %s", lm.Msg)
+}
+
+func TestCustomLoggingFormatter(t *testing.T) {
+	// beego.BConfig.Log.AccessLogs = true
+
+	logs.SetLoggerWithOpts("console", []string{`{"color":true}`}, common.SimpleKV{Key: "formatter", Value: customFormatter})
+
+	// Message will be formatted by the customFormatter with colorful text set to true
+	logs.Informational("Test message")
+}
+
+func TestGlobalLoggingFormatter(t *testing.T) {
+	logs.SetGlobalFormatter(globalFormatter)
+
+	logs.SetLogger("console", `{"color":true}`)
+
+	// Message will be formatted by globalFormatter
+	logs.Informational("Test message")
+
+}