Add conditional context and color on default formatter (#75)

* Add conditional context and color on default formatter * fix test
gol4ng · Mar 9, 2021 · 187d7c4 · 187d7c4
1 parent 13ed500
commit 187d7c4
Show file tree

Hide file tree

Showing 2 changed files with 84 additions and 7 deletions.
diff --git a/formatter/default.go b/formatter/default.go
@@ -6,17 +6,23 @@ import (
 	"github.com/gol4ng/logger"
 )
 
+var falseCondition = func(entry logger.Entry) bool {
+	return false
+}
+
 // DefaultFormatter is the default Entry formatter
 type DefaultFormatter struct {
-	colored        bool
-	displayContext bool
+	colored        func(entry logger.Entry) bool
+	displayContext func(entry logger.Entry) bool
 }
 
 // Format will return Entry as string
 func (n *DefaultFormatter) Format(entry logger.Entry) string {
+	n.init()
 	builder := &strings.Builder{}
 
-	if n.colored {
+	colored := n.colored(entry)
+	if colored {
 		switch entry.Level {
 		case logger.DebugLevel:
 			builder.WriteString("\x1b[1;36m")
@@ -40,20 +46,29 @@ func (n *DefaultFormatter) Format(entry logger.Entry) string {
 	builder.WriteString("<")
 	builder.WriteString(entry.Level.String())
 	builder.WriteString(">")
-	if n.colored {
+	if colored {
 		builder.WriteString("\x1b[m")
 	}
 	if entry.Message != "" {
 		builder.WriteString(" ")
 		builder.WriteString(entry.Message)
 	}
-	if n.displayContext && entry.Context != nil {
+	if entry.Context != nil && n.displayContext(entry) {
 		builder.WriteString(" ")
 		ContextToJSON(entry.Context, builder)
 	}
 	return builder.String()
 }
 
+func (n *DefaultFormatter) init() {
+	if n.colored == nil {
+		n.colored = falseCondition
+	}
+	if n.displayContext == nil {
+		n.displayContext = falseCondition
+	}
+}
+
 // NewDefaultFormatter will create a new DefaultFormatter
 func NewDefaultFormatter(options ...Option) *DefaultFormatter {
 	f := &DefaultFormatter{}
@@ -68,14 +83,28 @@ type Option func(*DefaultFormatter)
 
 // WithColor function will enable ANSI colored formatting
 func WithColor(enable bool) Option {
+	return WithConditionalColor(func(_ logger.Entry) bool {
+		return enable
+	})
+}
+
+// WithConditionalColor function will enable ANSI colored formatting
+func WithConditionalColor(conditional func(_ logger.Entry) bool) Option {
 	return func(formatter *DefaultFormatter) {
-		formatter.colored = enable
+		formatter.colored = conditional
 	}
 }
 
 // WithContext function will display context printing
 func WithContext(enable bool) Option {
+	return WithConditionalContext(func(_ logger.Entry) bool {
+		return enable
+	})
+}
+
+// WithConditionalContext function will display context printing
+func WithConditionalContext(conditional func(_ logger.Entry) bool) Option {
 	return func(formatter *DefaultFormatter) {
-		formatter.displayContext = enable
+		formatter.displayContext = conditional
 	}
 }
diff --git a/formatter/default_test.go b/formatter/default_test.go
@@ -68,6 +68,54 @@ func TestDefaultFormatter_Format_AllColor(t *testing.T) {
 	}
 }
 
+func TestDefaultFormatter_Format_ConditionalColor(t *testing.T) {
+	tests := []struct {
+		level    logger.Level
+		expected string
+	}{
+		{level: logger.EmergencyLevel, expected: "\x1b[1;37;41m<emergency>\x1b[m my message"},
+		{level: logger.AlertLevel, expected: "\x1b[1;30;43m<alert>\x1b[m my message"},
+		{level: logger.CriticalLevel, expected: "\x1b[1;30;47m<critical>\x1b[m my message"},
+		{level: logger.ErrorLevel, expected: "\x1b[1;31m<error>\x1b[m my message"},
+		{level: logger.WarningLevel, expected: "\x1b[1;33m<warning>\x1b[m my message"},
+		{level: logger.NoticeLevel, expected: "<notice> my message"},
+		{level: logger.InfoLevel, expected: "<info> my message"},
+		{level: logger.DebugLevel, expected: "<debug> my message"},
+	}
+	defaultFormatter := formatter.NewDefaultFormatter(formatter.WithConditionalColor(func(e logger.Entry) bool {
+		return e.Level <= logger.WarningLevel
+	}))
+	for _, tt := range tests {
+		t.Run(tt.level.String(), func(t *testing.T) {
+			assert.Equal(t, tt.expected, defaultFormatter.Format(logger.Entry{Level: tt.level, Message: "my message"}))
+		})
+	}
+}
+
+func TestDefaultFormatter_Format_ConditionalContext(t *testing.T) {
+	tests := []struct {
+		level    logger.Level
+		expected string
+	}{
+		{level: logger.EmergencyLevel, expected: `<emergency> my message {"my_name":"my value"}`},
+		{level: logger.AlertLevel, expected: `<alert> my message {"my_name":"my value"}`},
+		{level: logger.CriticalLevel, expected: `<critical> my message {"my_name":"my value"}`},
+		{level: logger.ErrorLevel, expected: `<error> my message {"my_name":"my value"}`},
+		{level: logger.WarningLevel, expected: `<warning> my message {"my_name":"my value"}`},
+		{level: logger.NoticeLevel, expected: "<notice> my message"},
+		{level: logger.InfoLevel, expected: "<info> my message"},
+		{level: logger.DebugLevel, expected: "<debug> my message"},
+	}
+	defaultFormatter := formatter.NewDefaultFormatter(formatter.WithConditionalContext(func(e logger.Entry) bool {
+		return e.Level <= logger.WarningLevel
+	}))
+	for _, tt := range tests {
+		t.Run(tt.level.String(), func(t *testing.T) {
+			assert.Equal(t, tt.expected, defaultFormatter.Format(logger.Entry{Level: tt.level, Message: "my message", Context: logger.Ctx("my_name", "my value")}))
+		})
+	}
+}
+
 // =====================================================================================================================
 // ================================================= EXAMPLES ==========================================================
 // =====================================================================================================================