Add a LogListener option to Notepad

For clients using the "log counting" feature, this will be slightly breaking. But it adds lots of flexibility which should make it worth while.

Now you can supply one or more `LogListener` funcs when creating a new `Notepad`. Typical use cases would be to capture ERROR logs in unit tests, and to count log statements. A listener for log counting is provided, see the `LogCounter` func.

If you want to add listeners to the global logger, see `SetLogListeners`.

.gitignore

@@ -20,3 +20,5 @@ _cgo_export.*
 _testmain.go
 
 *.exe
+*.bench
+go.sum

default_notepad.go

@@ -76,6 +76,13 @@ func SetFlags(flags int) {
 	reloadDefaultNotepad()
 }
 
+// SetLogListeners configures the default logger with one or more log listeners.
+func SetLogListeners(l ...LogListener) {
+	defaultNotepad.logListeners = l
+	defaultNotepad.init()
+	reloadDefaultNotepad()
+}
+
 // Level returns the current global log threshold.
 func LogThreshold() Threshold {
 	return defaultNotepad.logThreshold
@@ -95,19 +102,3 @@ func GetLogThreshold() Threshold {
 func GetStdoutThreshold() Threshold {
 	return defaultNotepad.GetStdoutThreshold()
 }
-
-// LogCountForLevel returns the number of log invocations for a given threshold.
-func LogCountForLevel(l Threshold) uint64 {
-	return defaultNotepad.LogCountForLevel(l)
-}
-
-// LogCountForLevelsGreaterThanorEqualTo returns the number of log invocations
-// greater than or equal to a given threshold.
-func LogCountForLevelsGreaterThanorEqualTo(threshold Threshold) uint64 {
-	return defaultNotepad.LogCountForLevelsGreaterThanorEqualTo(threshold)
-}
-
-// ResetLogCounters resets the invocation counters for all levels.
-func ResetLogCounters() {
-	defaultNotepad.ResetLogCounters()
-}

default_notepad_test.go

@@ -59,11 +59,16 @@ func TestDefaultLogging(t *testing.T) {
 }
 
 func TestLogCounter(t *testing.T) {
+	assert := require.New(t)
+
 	defaultNotepad.logHandle = ioutil.Discard
 	defaultNotepad.outHandle = ioutil.Discard
 
+	errorCounter := &Counter{}
+
 	SetLogThreshold(LevelTrace)
 	SetStdoutThreshold(LevelTrace)
+	SetLogListeners(LogCounter(errorCounter, LevelError))
 
 	FATAL.Println("fatal err")
 	CRITICAL.Println("critical err")
@@ -82,21 +87,13 @@ func TestLogCounter(t *testing.T) {
 			for j := 0; j < 10; j++ {
 				ERROR.Println("error", j)
 				// check for data races
-				require.True(t, LogCountForLevel(LevelError) > uint64(j))
-				require.True(t, LogCountForLevelsGreaterThanorEqualTo(LevelError) > uint64(j))
+				assert.True(errorCounter.Count() > uint64(j))
 			}
 		}()
 
 	}
 
 	wg.Wait()
 
-	require.Equal(t, uint64(1), LogCountForLevel(LevelFatal))
-	require.Equal(t, uint64(1), LogCountForLevel(LevelCritical))
-	require.Equal(t, uint64(2), LogCountForLevel(LevelWarn))
-	require.Equal(t, uint64(1), LogCountForLevel(LevelInfo))
-	require.Equal(t, uint64(1), LogCountForLevel(LevelDebug))
-	require.Equal(t, uint64(1), LogCountForLevel(LevelTrace))
-	require.Equal(t, uint64(100), LogCountForLevel(LevelError))
-	require.Equal(t, uint64(102), LogCountForLevelsGreaterThanorEqualTo(LevelError))
+	assert.Equal(uint64(102), errorCounter.Count())
 }

go.mod

@@ -1 +1,7 @@
 module github.com/spf13/jwalterweatherman
+
+require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/stretchr/testify v1.2.2
+)

log_counter.go

@@ -6,50 +6,41 @@
 package jwalterweatherman
 
 import (
+	"io"
 	"sync/atomic"
 )
 
-type logCounter struct {
-	counter uint64
+// Counter is an io.Writer that increments a counter on Write.
+type Counter struct {
+	count uint64
 }
 
-func (c *logCounter) incr() {
-	atomic.AddUint64(&c.counter, 1)
+func (c *Counter) incr() {
+	atomic.AddUint64(&c.count, 1)
 }
 
-func (c *logCounter) resetCounter() {
-	atomic.StoreUint64(&c.counter, 0)
+// Reset resets the counter.
+func (c *Counter) Reset() {
+	atomic.StoreUint64(&c.count, 0)
 }
 
-func (c *logCounter) getCount() uint64 {
-	return atomic.LoadUint64(&c.counter)
+// Count returns the current count.
+func (c *Counter) Count() uint64 {
+	return atomic.LoadUint64(&c.count)
 }
 
-func (c *logCounter) Write(p []byte) (n int, err error) {
+func (c *Counter) Write(p []byte) (n int, err error) {
 	c.incr()
 	return len(p), nil
 }
 
-// LogCountForLevel returns the number of log invocations for a given threshold.
-func (n *Notepad) LogCountForLevel(l Threshold) uint64 {
-	return n.logCounters[l].getCount()
-}
-
-// LogCountForLevelsGreaterThanorEqualTo returns the number of log invocations
-// greater than or equal to a given threshold.
-func (n *Notepad) LogCountForLevelsGreaterThanorEqualTo(threshold Threshold) uint64 {
-	var cnt uint64
-
-	for i := int(threshold); i < len(n.logCounters); i++ {
-		cnt += n.LogCountForLevel(Threshold(i))
-	}
-
-	return cnt
-}
-
-// ResetLogCounters resets the invocation counters for all levels.
-func (n *Notepad) ResetLogCounters() {
-	for _, np := range n.logCounters {
-		np.resetCounter()
+// LogCounter creates a LogListener that counts log statements >= the given threshold.
+func LogCounter(counter *Counter, t1 Threshold) LogListener {
+	return func(t2 Threshold) io.Writer {
+		if t2 < t1 {
+			// Not interested in this threshold.
+			return nil
+		}
+		return counter
 	}
 }

notepad.go

@@ -8,6 +8,7 @@ package jwalterweatherman
 import (
 	"fmt"
 	"io"
+	"io/ioutil"
 	"log"
 )
 
@@ -58,13 +59,28 @@ type Notepad struct {
 	prefix          string
 	flags           int
 
-	// One per Threshold
-	logCounters [7]*logCounter
+	logListeners []LogListener
 }
 
-// NewNotepad create a new notepad.
-func NewNotepad(outThreshold Threshold, logThreshold Threshold, outHandle, logHandle io.Writer, prefix string, flags int) *Notepad {
-	n := &Notepad{}
+// A LogListener can ble supplied to a Notepad to listen on log writes for a given
+// threshold. This can be used to capture log events in unit tests and similar.
+// Note that this function will be invoked once for each log threshold. If
+// the given threshold is not of interest to you, return nil.
+// Note that these listeners will receive log events for a given threshold, even
+// if the current configuration says not to log it. That way you can count ERRORs even
+// if you don't print them to the console.
+type LogListener func(t Threshold) io.Writer
+
+// NewNotepad creates a new Notepad.
+func NewNotepad(
+	outThreshold Threshold,
+	logThreshold Threshold,
+	outHandle, logHandle io.Writer,
+	prefix string, flags int,
+	logListeners ...LogListener,
+) *Notepad {
+
+	n := &Notepad{logListeners: logListeners}
 
 	n.loggers = [7]**log.Logger{&n.TRACE, &n.DEBUG, &n.INFO, &n.WARN, &n.ERROR, &n.CRITICAL, &n.FATAL}
 	n.outHandle = outHandle
@@ -95,26 +111,41 @@ func (n *Notepad) init() {
 
 	for t, logger := range n.loggers {
 		threshold := Threshold(t)
-		counter := &logCounter{}
-		n.logCounters[t] = counter
 		prefix := n.prefix + threshold.String() + " "
 
 		switch {
 		case threshold >= n.logThreshold && threshold >= n.stdoutThreshold:
-			*logger = log.New(io.MultiWriter(counter, logAndOut), prefix, n.flags)
+			*logger = log.New(n.createLogWriters(threshold, logAndOut), prefix, n.flags)
 
 		case threshold >= n.logThreshold:
-			*logger = log.New(io.MultiWriter(counter, n.logHandle), prefix, n.flags)
+			*logger = log.New(n.createLogWriters(threshold, n.logHandle), prefix, n.flags)
 
 		case threshold >= n.stdoutThreshold:
-			*logger = log.New(io.MultiWriter(counter, n.outHandle), prefix, n.flags)
+			*logger = log.New(n.createLogWriters(threshold, n.outHandle), prefix, n.flags)
 
 		default:
-			// counter doesn't care about prefix and flags, so don't use them
-			// for performance.
-			*logger = log.New(counter, "", 0)
+			*logger = log.New(n.createLogWriters(threshold, ioutil.Discard), prefix, n.flags)
+		}
+	}
+}
+
+func (n *Notepad) createLogWriters(t Threshold, handle io.Writer) io.Writer {
+	if len(n.logListeners) == 0 {
+		return handle
+	}
+	writers := []io.Writer{handle}
+	for _, l := range n.logListeners {
+		w := l(t)
+		if w != nil {
+			writers = append(writers, w)
 		}
 	}
+
+	if len(writers) == 1 {
+		return handle
+	}
+
+	return io.MultiWriter(writers...)
 }
 
 // SetLogThreshold changes the threshold above which messages are written to the

notepad_test.go

@@ -7,6 +7,8 @@ package jwalterweatherman
 
 import (
 	"bytes"
+	"io"
+	"io/ioutil"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -15,7 +17,9 @@ import (
 func TestNotepad(t *testing.T) {
 	var logHandle, outHandle bytes.Buffer
 
-	n := NewNotepad(LevelCritical, LevelError, &outHandle, &logHandle, "TestNotePad", 0)
+	errorCounter := &Counter{}
+
+	n := NewNotepad(LevelCritical, LevelError, &outHandle, &logHandle, "TestNotePad", 0, LogCounter(errorCounter, LevelError))
 
 	require.Equal(t, LevelCritical, n.GetStdoutThreshold())
 	require.Equal(t, LevelError, n.GetLogThreshold())
@@ -29,9 +33,48 @@ func TestNotepad(t *testing.T) {
 	require.NotContains(t, outHandle.String(), "Some error")
 	require.Contains(t, outHandle.String(), "Some critical error")
 
-	require.Equal(t, n.LogCountForLevel(LevelError), uint64(1))
-	require.Equal(t, n.LogCountForLevel(LevelDebug), uint64(1))
-	require.Equal(t, n.LogCountForLevel(LevelTrace), uint64(0))
+	// 1 error + 1 critical
+	require.Equal(t, errorCounter.Count(), uint64(2))
+}
+
+func TestNotepadLogListener(t *testing.T) {
+	assert := require.New(t)
+
+	var errorBuff, infoBuff bytes.Buffer
+
+	errorCapture := func(t Threshold) io.Writer {
+		if t != LevelError {
+			// Only interested in ERROR
+			return nil
+		}
+
+		return &errorBuff
+	}
+
+	infoCapture := func(t Threshold) io.Writer {
+		if t != LevelInfo {
+			return nil
+		}
+
+		return &infoBuff
+	}
+
+	n := NewNotepad(LevelCritical, LevelError, ioutil.Discard, ioutil.Discard, "TestNotePad", 0, infoCapture, errorCapture)
+
+	n.DEBUG.Println("Some debug")
+	n.INFO.Println("Some info")
+	n.INFO.Println("Some more info")
+	n.ERROR.Println("Some error")
+	n.CRITICAL.Println("Some critical error")
+	n.ERROR.Println("Some more error")
+
+	assert.Equal(`[TestNotePad] ERROR Some error
+[TestNotePad] ERROR Some more error
+`, errorBuff.String())
+	assert.Equal(`[TestNotePad] INFO Some info
+[TestNotePad] INFO Some more info
+`, infoBuff.String())
+
 }
 
 func TestThresholdString(t *testing.T) {