From e6288e90eb313122410bf1335d502a062c1ae610 Mon Sep 17 00:00:00 2001
From: Sai <sairoutine@gmail.com>
Date: Fri, 22 Feb 2019 17:48:55 +0900
Subject: [PATCH] Change color methods in using defaultLogger function to
 public (#1771)

Fix https://github.com/gin-gonic/gin/issues/1768
---
 logger.go      | 85 ++++++++++++++++++++++++++++----------------------
 logger_test.go | 19 +++++++++++
 2 files changed, 67 insertions(+), 37 deletions(-)

diff --git a/logger.go b/logger.go
index bd28a11..dc63997 100644
--- a/logger.go
+++ b/logger.go
@@ -68,13 +68,58 @@ type LogFormatterParams struct {
 	BodySize int
 }
 
+// StatusCodeColor is the ANSI color for appropriately logging http status code to a terminal.
+func (p *LogFormatterParams) StatusCodeColor() string {
+	code := p.StatusCode
+
+	switch {
+	case code >= http.StatusOK && code < http.StatusMultipleChoices:
+		return green
+	case code >= http.StatusMultipleChoices && code < http.StatusBadRequest:
+		return white
+	case code >= http.StatusBadRequest && code < http.StatusInternalServerError:
+		return yellow
+	default:
+		return red
+	}
+}
+
+// MethodColor is the ANSI color for appropriately logging http method to a terminal.
+func (p *LogFormatterParams) MethodColor() string {
+	method := p.Method
+
+	switch method {
+	case "GET":
+		return blue
+	case "POST":
+		return cyan
+	case "PUT":
+		return yellow
+	case "DELETE":
+		return red
+	case "PATCH":
+		return green
+	case "HEAD":
+		return magenta
+	case "OPTIONS":
+		return white
+	default:
+		return reset
+	}
+}
+
+// ResetColor resets all escape attributes.
+func (p *LogFormatterParams) ResetColor() string {
+	return reset
+}
+
 // defaultLogFormatter is the default log format function Logger middleware uses.
 var defaultLogFormatter = func(param LogFormatterParams) string {
 	var statusColor, methodColor, resetColor string
 	if param.IsTerm {
-		statusColor = colorForStatus(param.StatusCode)
-		methodColor = colorForMethod(param.Method)
-		resetColor = reset
+		statusColor = param.StatusCodeColor()
+		methodColor = param.MethodColor()
+		resetColor = param.ResetColor()
 	}
 
 	return fmt.Sprintf("[GIN] %v |%s %3d %s| %13v | %15s |%s %-7s %s %s\n%s",
@@ -205,37 +250,3 @@ func LoggerWithConfig(conf LoggerConfig) HandlerFunc {
 		}
 	}
 }
-
-func colorForStatus(code int) string {
-	switch {
-	case code >= http.StatusOK && code < http.StatusMultipleChoices:
-		return green
-	case code >= http.StatusMultipleChoices && code < http.StatusBadRequest:
-		return white
-	case code >= http.StatusBadRequest && code < http.StatusInternalServerError:
-		return yellow
-	default:
-		return red
-	}
-}
-
-func colorForMethod(method string) string {
-	switch method {
-	case "GET":
-		return blue
-	case "POST":
-		return cyan
-	case "PUT":
-		return yellow
-	case "DELETE":
-		return red
-	case "PATCH":
-		return green
-	case "HEAD":
-		return magenta
-	case "OPTIONS":
-		return white
-	default:
-		return reset
-	}
-}
diff --git a/logger_test.go b/logger_test.go
index 8770b5f..c551677 100644
--- a/logger_test.go
+++ b/logger_test.go
@@ -257,6 +257,13 @@ func TestDefaultLogFormatter(t *testing.T) {
 }
 
 func TestColorForMethod(t *testing.T) {
+	colorForMethod := func(method string) string {
+		p := LogFormatterParams{
+			Method: method,
+		}
+		return p.MethodColor()
+	}
+
 	assert.Equal(t, string([]byte{27, 91, 57, 55, 59, 52, 52, 109}), colorForMethod("GET"), "get should be blue")
 	assert.Equal(t, string([]byte{27, 91, 57, 55, 59, 52, 54, 109}), colorForMethod("POST"), "post should be cyan")
 	assert.Equal(t, string([]byte{27, 91, 57, 48, 59, 52, 51, 109}), colorForMethod("PUT"), "put should be yellow")
@@ -268,12 +275,24 @@ func TestColorForMethod(t *testing.T) {
 }
 
 func TestColorForStatus(t *testing.T) {
+	colorForStatus := func(code int) string {
+		p := LogFormatterParams{
+			StatusCode: code,
+		}
+		return p.StatusCodeColor()
+	}
+
 	assert.Equal(t, string([]byte{27, 91, 57, 55, 59, 52, 50, 109}), colorForStatus(http.StatusOK), "2xx should be green")
 	assert.Equal(t, string([]byte{27, 91, 57, 48, 59, 52, 55, 109}), colorForStatus(http.StatusMovedPermanently), "3xx should be white")
 	assert.Equal(t, string([]byte{27, 91, 57, 48, 59, 52, 51, 109}), colorForStatus(http.StatusNotFound), "4xx should be yellow")
 	assert.Equal(t, string([]byte{27, 91, 57, 55, 59, 52, 49, 109}), colorForStatus(2), "other things should be red")
 }
 
+func TestResetColor(t *testing.T) {
+	p := LogFormatterParams{}
+	assert.Equal(t, string([]byte{27, 91, 48, 109}), p.ResetColor())
+}
+
 func TestErrorLogger(t *testing.T) {
 	router := New()
 	router.Use(ErrorLogger())