2 files changed, 101 insertions, 0 deletions
diff --git a/logger/logger.go b/logger/logger.go
index bb76657..9cd4b33 100644
--- a/logger/logger.go
+++ b/logger/logger.go
@@ -3,6 +3,7 @@ package logger
 
 import (
 	"fmt"
+	"io"
 	"log"
 	"os"
 	"runtime/debug"
@@ -13,6 +14,7 @@ type severity int
 
 type logger interface {
 	Output(calldepth int, s string) error
+	SetOutput(w io.Writer)
 	SetFlags(flag int)
 }
 
@@ -103,6 +105,13 @@ func (l *Logger) output(s severity, depth int, txt string) {
 	l.loggers[s].Output(3+depth, txt+"\033[0m")
 }
 
+// SetOutput changes the output of the logger.
+func (l *Logger) SetOutput(w io.Writer) {
+	for _, logger := range l.loggers {
+		logger.SetOutput(w)
+	}
+}
+
 // SetFlags sets the output flags for the logger.
 func (l *Logger) SetFlags(flag int) {
 	for _, logger := range l.loggers {
@@ -194,6 +203,11 @@ func (l *Logger) Fatalf(format string, v ...interface{}) {
 	os.Exit(1)
 }
 
+// SetOutput changes the output of the default logger.
+func SetOutput(w io.Writer) {
+	defaultLogger.SetOutput(w)
+}
+
 // SetFlags sets the output flags for the logger.
 func SetFlags(flag int) {
 	defaultLogger.SetFlags(flag)
diff --git a/logger/logger_test.go b/logger/logger_test.go
new file mode 100644
index 0000000..ce708bb
--- /dev/null
+++ b/logger/logger_test.go
@@ -0,0 +1,87 @@
+package logger
+
+import (
+	"bufio"
+	"bytes"
+	"os"
+	"reflect"
+	"strings"
+	"testing"
+)
+
+func TestLoggingBeforeInit(t *testing.T) {
+	old := os.Stderr
+	r, w, err := os.Pipe()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	os.Stderr = w
+	// Reset
+	initialize()
+
+	info := "info log"
+	warning := "warning log"
+	errL := "error log"
+	fatal := "fatal log"
+
+	Info(info)
+	Warning(warning)
+	Error(errL)
+	// We don't want os.Exit in a test
+	defaultLogger.output(sFatal, 0, fatal)
+
+	w.Close()
+	os.Stderr = old
+
+	var b bytes.Buffer
+	scanner := bufio.NewScanner(r)
+	for scanner.Scan() {
+		b.Write(scanner.Bytes())
+	}
+
+	out := b.String()
+
+	for _, txt := range []string{info, warning, errL, fatal} {
+		if !strings.Contains(out, txt) {
+			t.Errorf("log output %q does not contain expected text: %q", out, txt)
+		}
+	}
+}
+
+func TestInit(t *testing.T) {
+	var buf1 bytes.Buffer
+	l1 := Init(3)
+	l1.SetOutput(&buf1)
+	if !reflect.DeepEqual(l1, defaultLogger) {
+		t.Fatal("defaultLogger does not match logger returned by Init")
+	}
+
+	// Subsequent runs of Init shouldn't change defaultLogger.
+	var buf2 bytes.Buffer
+	l2 := Init(3)
+	l2.SetOutput(&buf2)
+	if !reflect.DeepEqual(l1, defaultLogger) {
+		t.Error("defaultLogger should not have changed")
+	}
+
+	// Check log output.
+	l1.Info("logger #1")
+	l2.Info("logger #2")
+	defaultLogger.Info("default logger")
+
+	tests := []struct {
+		out  string
+		want int
+	}{
+		{buf1.String(), 2},
+		{buf2.String(), 1},
+	}
+
+	for i, tt := range tests {
+		got := len(strings.Split(strings.TrimSpace(tt.out), "\n"))
+		if got != tt.want {
+			t.Errorf("logger %d wrong number of lines, want %d, got %d", i+1, tt.want, got)
+		}
+	}
+}