Merge pull request #23 from Bose/banner

Banner

go.mod

@@ -5,6 +5,7 @@ require (
 	github.com/gin-gonic/gin v1.4.0
 	github.com/google/uuid v1.1.0
 	github.com/kr/pretty v0.1.0 // indirect
+	github.com/matryer/is v1.2.0
 	github.com/mitchellh/copystructure v1.0.0
 	github.com/opentracing/opentracing-go v1.1.0
 	github.com/prometheus/client_model v0.0.0-20190129233127-fd36f4220a90 // indirect

go.sum

@@ -40,6 +40,8 @@ github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORN
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
 github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
+github.com/matryer/is v1.2.0 h1:92UTHpy8CDwaJ08GqLDzhhuixiBUUD1p3AU6PHddz4A=
+github.com/matryer/is v1.2.0/go.mod h1:2fLPjFQM9rhQ15aVEtbuwhJinnOqrmgXPNdZsdwlWXA=
 github.com/mattn/go-isatty v0.0.4 h1:bnP0vzxcAdeI1zdubAl5PjU6zsERjGZb7raWodagDYs=
 github.com/mattn/go-isatty v0.0.4/go.mod h1:M+lRXTBqGeGNdLjl/ufCoiOlB5xdOkqRJdNxMWT7Zi4=
 github.com/mattn/go-isatty v0.0.7 h1:UvyT9uN+3r7yLEYSlJsbQGdsaB/a0DlgWP3pql6iwOc=

logBuffer.go

@@ -16,21 +16,24 @@ type LogBuffer struct {
 	header    map[string]interface{}
 	headerMU  *sync.RWMutex
 	AddBanner bool
+	banner    string
 	MaxSize   uint
 }
 
 // NewLogBuffer - create a LogBuffer and initialize it
 func NewLogBuffer(opt ...LogBufferOption) LogBuffer {
-	opts := logBufferOptions{maxSize: DefaultLogBufferMaxSize}
+	opts := defaultLogBufferOptions()
 	for _, o := range opt {
 		o(&opts)
 	}
-	return LogBuffer{
+	b := LogBuffer{
 		header:    opts.withHeaders,
 		headerMU:  &sync.RWMutex{},
 		AddBanner: opts.addBanner,
 		MaxSize:   opts.maxSize,
 	}
+	b.SetCustomBanner(opts.banner)
+	return b
 }
 
 // StoreHeader - store a header
@@ -118,8 +121,13 @@ func (b *LogBuffer) String() string {
 	}
 	str.WriteString("\"entries\":[" + strings.TrimSuffix(b.Buff.String(), ",") + "]")
 	if b.AddBanner {
-		str.WriteString(",\"banner\":\"[GIN] --------------------------------------------------------------- GinLogrusWithTracing ----------------------------------------------------------------\"")
+		str.WriteString(b.banner)
 	}
 	str.WriteString("}\n")
 	return str.String()
 }
+
+// SetCustomBanner allows a custom banner to be set after the NewLogBuffer() has been used
+func (b *LogBuffer) SetCustomBanner(banner string) {
+	b.banner = fmt.Sprintf(",\"banner\":\"%s\"", banner)
+}

logBufferOptions.go

@@ -1,18 +1,25 @@
 package ginlogrus
 
+const DefaultBanner = "[GIN] --------------------------------------------------------------- GinLogrusWithTracing ----------------------------------------------------------------"
+
 // LogBufferOption - define options for LogBuffer
 type LogBufferOption func(*logBufferOptions)
 type logBufferOptions struct {
 	addBanner   bool
 	withHeaders map[string]interface{}
 	maxSize     uint
+	banner      string
 }
 
 // DefaultLogBufferMaxSize - avg single spaced page contains 3k chars, so 100k == 33 pages which is a reasonable max
 const DefaultLogBufferMaxSize = 100000
 
 func defaultLogBufferOptions() logBufferOptions {
-	return logBufferOptions{maxSize: DefaultLogBufferMaxSize}
+	return logBufferOptions{
+		maxSize:   DefaultLogBufferMaxSize,
+		banner:    DefaultBanner,
+		addBanner: false,
+	}
 }
 
 // WithBanner - define an Option func for passing in an optional add Banner
@@ -31,8 +38,17 @@ func WithHeader(k string, v interface{}) LogBufferOption {
 		o.withHeaders[k] = v
 	}
 }
+
+// WithMaxSize specifies the bounded max size the logBuffer can grow to
 func WithMaxSize(s uint) LogBufferOption {
 	return func(o *logBufferOptions) {
 		o.maxSize = s
 	}
 }
+
+// WithCustomBanner allows users to define their own custom banner
+func WithCustomBanner(b string) LogBufferOption {
+	return func(o *logBufferOptions) {
+		o.banner = b
+	}
+}

logBufferOptions_test.go

@@ -24,6 +24,26 @@ func TestWithBanner(t *testing.T) {
 	}
 }
 
+func TestWithCustomBanner(t *testing.T) {
+	tests := []struct {
+		name string
+		want string
+	}{
+		{name: "custom", want: "CustomBanner"},
+		{name: "default", want: DefaultBanner},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			opts := defaultLogBufferOptions()
+			f := WithCustomBanner(tt.want)
+			f(&opts)
+			if opts.banner != tt.want {
+				t.Errorf("WithBanner() = %v, want %v", opts.addBanner, tt.want)
+			}
+		})
+	}
+}
+
 func TestWithHeader(t *testing.T) {
 	tests := []struct {
 		name      string

logBuffer_test.go

@@ -52,12 +52,22 @@ func TestNewLogBuffer(t *testing.T) {
 		{
 			name: "one",
 			opt:  []LogBufferOption{WithBanner(true), WithHeader("1", true)},
-			want: LogBuffer{AddBanner: true, header: map[string]interface{}{"1": true}, headerMU: &sync.RWMutex{}, MaxSize: DefaultLogBufferMaxSize},
+			want: LogBuffer{AddBanner: true, header: map[string]interface{}{"1": true}, headerMU: &sync.RWMutex{}, MaxSize: DefaultLogBufferMaxSize, banner: ",\"banner\":\"[GIN] --------------------------------------------------------------- GinLogrusWithTracing ----------------------------------------------------------------\""},
 		},
 		{
 			name: "two",
 			opt:  []LogBufferOption{WithHeader("1", "one"), WithHeader("2", true)},
-			want: LogBuffer{AddBanner: false, header: map[string]interface{}{"1": "one", "2": true}, headerMU: &sync.RWMutex{}, MaxSize: DefaultLogBufferMaxSize},
+			want: LogBuffer{AddBanner: false, header: map[string]interface{}{"1": "one", "2": true}, headerMU: &sync.RWMutex{}, MaxSize: DefaultLogBufferMaxSize, banner: ",\"banner\":\"[GIN] --------------------------------------------------------------- GinLogrusWithTracing ----------------------------------------------------------------\""},
+		},
+		{
+			name: "three",
+			opt:  []LogBufferOption{WithBanner(true), WithHeader("1", true), WithCustomBanner("custom")},
+			want: LogBuffer{AddBanner: true, header: map[string]interface{}{"1": true}, headerMU: &sync.RWMutex{}, MaxSize: DefaultLogBufferMaxSize, banner: ",\"banner\":\"custom\""},
+		},
+		{
+			name: "four",
+			opt:  []LogBufferOption{WithBanner(false), WithHeader("1", true)},
+			want: LogBuffer{AddBanner: false, header: map[string]interface{}{"1": true}, headerMU: &sync.RWMutex{}, MaxSize: DefaultLogBufferMaxSize, banner: ",\"banner\":\"[GIN] --------------------------------------------------------------- GinLogrusWithTracing ----------------------------------------------------------------\""},
 		},
 	}
 	for _, tt := range tests {

middleware.go

@@ -49,7 +49,7 @@ func WithTracing(
 	return func(c *gin.Context) {
 		// var aggregateLoggingBuff strings.Builder
 		// var aggregateLoggingBuff logBuffer
-		aggregateLoggingBuff := NewLogBuffer(WithBanner(true))
+		aggregateLoggingBuff := NewLogBuffer(WithBanner(useBanner), WithCustomBanner(opts.banner))
 		aggregateRequestLogger := &logrus.Logger{
 			Out:       &aggregateLoggingBuff,
 			Formatter: new(logrus.JSONFormatter),
@@ -121,7 +121,7 @@ func WithTracing(
 				// if useBanner {
 				// 	fields["banner"] = "[GIN] --------------------------------------------------------------- GinLogrusWithTracing ----------------------------------------------------------------"
 				// }
-				fmt.Printf(aggregateLoggingBuff.String())
+				fmt.Fprintf(opts.writer, aggregateLoggingBuff.String())
 			}
 		}
 	}

middleware_test.go

@@ -0,0 +1,91 @@
+package ginlogrus
+
+import (
+	"bytes"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"github.com/matryer/is"
+	"github.com/sirupsen/logrus"
+)
+
+func performRequest(method, target string, router *gin.Engine) *httptest.ResponseRecorder {
+	r := httptest.NewRequest(method, target, nil)
+	w := httptest.NewRecorder()
+	router.ServeHTTP(w, r)
+	return w
+}
+
+func TestBanner(t *testing.T) {
+	is := is.New(t)
+	buff := ""
+	getHandler := func(c *gin.Context) {
+		SetCtxLoggerHeader(c, "new-header-index-name", "this is how you set new header level data")
+
+		logger := GetCtxLogger(c)
+		logger.Info("test-entry-1")
+		logger.Info("test-entry-2")
+		c.JSON(200, "Hello world!")
+	}
+	gin.SetMode(gin.DebugMode)
+	gin.DisableConsoleColor()
+
+	l := bytes.NewBufferString(buff)
+	r := gin.Default()
+	r.Use(WithTracing(logrus.StandardLogger(),
+		false,
+		time.RFC3339,
+		true,
+		"requestID",
+		[]byte("uber-trace-id"), // where jaeger might have put the trace id
+		[]byte("RequestID"),     // where the trace ID might already be populated in the headers
+		WithAggregateLogging(true),
+		WithWriter(l)))
+	r.GET("/", getHandler)
+	w := performRequest("GET", "/", r)
+	is.Equal(200, w.Code)
+	t.Log("this is the buffer: ", l)
+	is.True(!strings.Contains(l.String(), "GinLogrusWithTracing"))
+
+	buff = ""
+	l = bytes.NewBufferString(buff)
+	r = gin.New()
+	r.Use(WithTracing(logrus.StandardLogger(),
+		true,
+		time.RFC3339,
+		true,
+		"requestID",
+		[]byte("uber-trace-id"), // where jaeger might have put the trace id
+		[]byte("RequestID"),     // where the trace ID might already be populated in the headers
+		WithAggregateLogging(true),
+		WithWriter(l)))
+	r.GET("/", getHandler)
+	w = performRequest("GET", "/", r)
+	is.Equal(200, w.Code)
+	t.Log("this is the buffer: ", l)
+	is.True(strings.Contains(l.String(), "GinLogrusWithTracing"))
+
+	customBanner := "---- custom banner ----"
+	buff = ""
+	l = bytes.NewBufferString(buff)
+	r = gin.New()
+	r.Use(WithTracing(logrus.StandardLogger(),
+		true,
+		time.RFC3339,
+		true,
+		"requestID",
+		[]byte("uber-trace-id"), // where jaeger might have put the trace id
+		[]byte("RequestID"),     // where the trace ID might already be populated in the headers
+		WithLogCustomBanner(customBanner),
+		WithAggregateLogging(true),
+		WithWriter(l)))
+	r.GET("/", getHandler)
+	w = performRequest("GET", "/", r)
+	is.Equal(200, w.Code)
+	t.Log("this is the buffer: ", l)
+	is.True(strings.Contains(l.String(), customBanner))
+
+}

options.go

@@ -1,18 +1,27 @@
 package ginlogrus
 
-import "github.com/sirupsen/logrus"
+import (
+	"io"
+	"os"
+
+	"github.com/sirupsen/logrus"
+)
 
 // Option - define options for WithTracing()
 type Option func(*options)
 type options struct {
 	aggregateLogging bool
-	logLevel logrus.Level
+	logLevel         logrus.Level
+	writer           io.Writer
+	banner           string
 }
 
 // defaultOptions - some defs options to NewJWTCache()
 var defaultOptions = options{
 	aggregateLogging: false,
-	logLevel: logrus.DebugLevel,
+	logLevel:         logrus.DebugLevel,
+	writer:           os.Stdout,
+	banner:           DefaultBanner,
 }
 
 // WithAggregateLogging - define an Option func for passing in an optional aggregateLogging
@@ -28,3 +37,18 @@ func WithLogLevel(logLevel logrus.Level) Option {
 		o.logLevel = logLevel
 	}
 }
+
+// WithWriter allows users to define the writer used for middlware aggregagte logging, the default writer is os.Stdout
+func WithWriter(w io.Writer) Option {
+	return func(o *options) {
+		o.writer = w
+	}
+}
+
+// WithLogCustomBanner allows users to define their own custom banner.  There is some overlap with this name and the LogBufferOption.CustomBanner and yes,
+// they are related options, but I didn't want to make a breaking API change to support this new option... so we'll have to live with a bit of confusion/overlap in option names
+func WithLogCustomBanner(b string) Option {
+	return func(o *options) {
+		o.banner = b
+	}
+}