New TestWriter for logging to testing.TB (#369)

writer.go

@@ -1,7 +1,12 @@
 package zerolog
 
 import (
+	"bytes"
 	"io"
+	"path"
+	"runtime"
+	"strconv"
+	"strings"
 	"sync"
 )
 
@@ -96,3 +101,54 @@ func MultiLevelWriter(writers ...io.Writer) LevelWriter {
 	}
 	return multiLevelWriter{lwriters}
 }
+
+// TestingLog is the logging interface of testing.TB.
+type TestingLog interface {
+	Log(args ...interface{})
+	Logf(format string, args ...interface{})
+	Helper()
+}
+
+// TestWriter is a writer that writes to testing.TB.
+type TestWriter struct {
+	T TestingLog
+
+	// Frame skips caller frames to capture the original file and line numbers.
+	Frame int
+}
+
+// NewTestWriter creates a writer that logs to the testing.TB.
+func NewTestWriter(t TestingLog) TestWriter {
+	return TestWriter{T: t}
+}
+
+// Write to testing.TB.
+func (t TestWriter) Write(p []byte) (n int, err error) {
+	t.T.Helper()
+
+	n = len(p)
+
+	// Strip trailing newline because t.Log always adds one.
+	p = bytes.TrimRight(p, "\n")
+
+	// Try to correct the log file and line number to the caller.
+	if t.Frame > 0 {
+		_, origFile, origLine, _ := runtime.Caller(1)
+		_, frameFile, frameLine, ok := runtime.Caller(1 + t.Frame)
+		if ok {
+			erase := strings.Repeat("\b", len(path.Base(origFile))+len(strconv.Itoa(origLine))+3)
+			t.T.Logf("%s%s:%d: %s", erase, path.Base(frameFile), frameLine, p)
+			return n, err
+		}
+	}
+	t.T.Log(string(p))
+
+	return n, err
+}
+
+// ConsoleTestWriter creates an option that correctly sets the file frame depth for testing.TB log.
+func ConsoleTestWriter(t TestingLog) func(w *ConsoleWriter) {
+	return func(w *ConsoleWriter) {
+		w.Out = TestWriter{T: t, Frame: 6}
+	}
+}

writer_test.go

@@ -1,10 +1,12 @@
-// +build !binary_log
-// +build !windows
+//go:build !binary_log && !windows
+// +build !binary_log,!windows
 
 package zerolog
 
 import (
+	"bytes"
 	"errors"
+	"fmt"
 	"io"
 	"reflect"
 	"testing"
@@ -111,3 +113,65 @@ func TestResilientMultiWriter(t *testing.T) {
 		writeCalls = 0
 	}
 }
+
+type testingLog struct {
+	testing.TB
+	buf bytes.Buffer
+}
+
+func (t *testingLog) Log(args ...interface{}) {
+	if _, err := t.buf.WriteString(fmt.Sprint(args...)); err != nil {
+		t.Error(err)
+	}
+}
+
+func (t *testingLog) Logf(format string, args ...interface{}) {
+	if _, err := t.buf.WriteString(fmt.Sprintf(format, args...)); err != nil {
+		t.Error(err)
+	}
+}
+
+func TestTestWriter(t *testing.T) {
+	tests := []struct {
+		name  string
+		write []byte
+		want  []byte
+	}{{
+		name:  "newline",
+		write: []byte("newline\n"),
+		want:  []byte("newline"),
+	}, {
+		name:  "oneline",
+		write: []byte("oneline"),
+		want:  []byte("oneline"),
+	}, {
+		name:  "twoline",
+		write: []byte("twoline\n\n"),
+		want:  []byte("twoline"),
+	}}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tb := &testingLog{TB: t} // Capture TB log buffer.
+			w := TestWriter{T: tb}
+
+			n, err := w.Write(tt.write)
+			if err != nil {
+				t.Error(err)
+			}
+			if n != len(tt.write) {
+				t.Errorf("Expected %d write length but got %d", len(tt.write), n)
+			}
+			p := tb.buf.Bytes()
+			if !bytes.Equal(tt.want, p) {
+				t.Errorf("Expected %q, got %q.", tt.want, p)
+			}
+
+			log := New(NewConsoleWriter(ConsoleTestWriter(t)))
+			log.Info().Str("name", tt.name).Msg("Success!")
+
+			tb.buf.Reset()
+		})
+	}
+
+}