Fix: improve test cases, implement zerolog.Hook for prefix

README.md

@@ -48,7 +48,7 @@ func main() {
 	*/
 	
 	// The zwrap solution, wrap the logger:
-	zl := zwrap.Wrap(&logger)
+	zl := zwrap.Wrap(logger)
 	myThing.SetLogger(zl)
 	myThing.DoSomething()
 }

wrap.go

@@ -68,9 +68,6 @@ type Logger struct {
 func (l *Logger) SetPrefix(prefix string) {
 	l.Lock()
 	l.prefix = prefix
-	l.Logger = nil
-	nl := l.Logger.With().Str("caller", prefix).Logger()
-	l.Logger = &nl
 	l.Unlock()
 }
 
@@ -219,9 +216,22 @@ func printLn(e *zerolog.Event, v ...interface{}) {
 	strBufs.Put(strBuf)
 }
 
-func Wrap(l *zerolog.Logger) *Logger {
-	return &Logger{
-		Logger:  l,
-		RWMutex: &sync.RWMutex{},
+type prefixHook struct {
+	parent StdCompatLogger
+}
+
+func (h prefixHook) Run(e *zerolog.Event, _ zerolog.Level, _ string) {
+	if h.parent.Prefix() != "" {
+		e.Str("caller", h.parent.Prefix())
 	}
 }
+
+func Wrap(l zerolog.Logger) *Logger {
+	wrapped := &Logger{
+		RWMutex: &sync.RWMutex{},
+	}
+	p := prefixHook{wrapped}
+	l = l.Hook(p)
+	wrapped.Logger = &l
+	return wrapped
+}

wrap_test.go

@@ -1,7 +1,9 @@
 package zwrap
 
 import (
+	"errors"
 	"log"
+	"os"
 	"strings"
 	"testing"
 
@@ -9,11 +11,25 @@ import (
 )
 
 type testWriter struct {
-	t *testing.T
+	t                 *testing.T
+	needsPrefix       string
+	mustNotHavePrefix string
 }
 
+var ErrPrefixMismatch = errors.New("prefix mismatch")
+
 func (w *testWriter) Write(p []byte) (n int, err error) {
-	w.t.Log(strings.TrimSuffix(string(p), "\n"))
+	w.t.Helper()
+	line := strings.TrimSuffix(string(p), "\n")
+	if w.needsPrefix != "" && !strings.Contains(line, w.needsPrefix) {
+		w.t.Errorf("expected prefix %q, got %q", w.needsPrefix, line)
+		return 0, ErrPrefixMismatch
+	}
+	if w.mustNotHavePrefix != "" && strings.Contains(line, w.mustNotHavePrefix) {
+		w.t.Errorf("unexpected prefix %q, got %q", w.mustNotHavePrefix, line)
+		return 0, ErrPrefixMismatch
+	}
+	w.t.Log(line)
 	return len(p), nil
 }
 
@@ -33,13 +49,64 @@ func (n *needsLogger) DoSomething() {
 	n.logger.Println("Hello, world!")
 }
 
-func TestWrapper(t *testing.T) {
+func TestWrap(t *testing.T) {
+	ExampleWrap()
+
+	writah := &testWriter{t: t}
+
+	zl := zerolog.New(writah).With().Timestamp().Logger()
+	wrapped := Wrap(zl)
+	myThing := &needsLogger{}
+	myThing.SetLogger(wrapped)
+
+	multiLog := func(v ...interface{}) {
+		t.Helper()
+		wrapped.Print(v...)
+		wrapped.Printf("%v", v)
+		wrapped.Println(v...)
+		wrapped.Error(v...)
+		wrapped.Errorf("%v", v)
+		wrapped.Errorln(v...)
+		wrapped.Debug(v...)
+		wrapped.Debugf("%v", v)
+		wrapped.Debugln(v...)
+		wrapped.Warn(v...)
+		wrapped.Warnf("%v", v)
+		wrapped.Warnln(v...)
+		wrapped.Info(v...)
+		wrapped.Infof("%v", v)
+		wrapped.Infoln(v...)
+		wrapped.Tracef("%v", v)
+		wrapped.Trace(v...)
+		wrapped.Traceln(v...)
+	}
+
+	t.Run("generic", func(t *testing.T) {
+		multiLog("Hello, world!")
+	})
+
+	t.Run("prefix", func(t *testing.T) {
+		writah.needsPrefix = "prefix: "
+		wrapped.SetPrefix("prefix: ")
+		multiLog("Hello, world!")
+	})
+
+	t.Run("remove prefix", func(t *testing.T) {
+		writah.needsPrefix = ""
+		writah.mustNotHavePrefix = "prefix: "
+		wrapped.SetPrefix("")
+		multiLog("Hello, world!")
+	})
+
+}
+
+func ExampleWrap() {
 	// Create a new zerolog.Logger
-	logger := zerolog.New(&testWriter{t: t}).With().Timestamp().Logger()
+	logger := zerolog.New(os.Stderr).With().Timestamp().Logger()
 
 	// Demonstrate that we can use the stdlib logger
 	myThing := &needsLogger{}
-	myThing.SetLogger(log.New(&testWriter{t: t}, "stdlog: ", log.LstdFlags))
+	myThing.SetLogger(log.New(os.Stderr, "stdlog: ", log.LstdFlags))
 	myThing.DoSomething()
 
 	// Demonstrate that we can use zerolog when wrapped
@@ -50,7 +117,7 @@ func TestWrapper(t *testing.T) {
 	*/
 
 	// The zwrap solution, wrap the logger:
-	zl := Wrap(&logger)
+	zl := Wrap(logger)
 	myThing.SetLogger(zl)
 	myThing.DoSomething()
 }