1 files changed, 44 insertions, 13 deletions
diff --git a/vendor/github.com/spf13/jwalterweatherman/notepad.go b/vendor/github.com/spf13/jwalterweatherman/notepad.go
index ae5aaf71..cc7957bf 100644
--- a/vendor/github.com/spf13/jwalterweatherman/notepad.go
+++ b/vendor/github.com/spf13/jwalterweatherman/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