New TestWriter for logging to testing.TB

writer.go

@@ -1,8 +1,14 @@
 package zerolog
 
 import (
+	"bytes"
 	"io"
+	"path"
+	"runtime"
+	"strconv"
+	"strings"
 	"sync"
+	"testing"
 )
 
 // LevelWriter defines as interface a writer may implement in order
@@ -96,3 +102,45 @@ func MultiLevelWriter(writers ...io.Writer) LevelWriter {
 	}
 	return multiLevelWriter{lwriters}
 }
+
+// TestWriter is a writer that writes to testing.TB.
+type TestWriter struct {
+	TB testing.TB
+
+	// 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 testing.TB) TestWriter {
+	return TestWriter{TB: t, Frame: 1}
+}
+
+// Write to testing.TB.
+func (t TestWriter) Write(p []byte) (n int, err error) {
+	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(0)
+		_, frameFile, frameLine, ok := runtime.Caller(t.Frame)
+		if ok {
+			erase := strings.Repeat("\b", len(path.Base(origFile))+len(strconv.Itoa(origLine))+3)
+			t.TB.Logf("%s%s:%d: %s", erase, path.Base(frameFile), frameLine, p)
+			return n, err
+		}
+	}
+	t.TB.Logf("%s", p)
+
+	return n, err
+}
+
+// ConsoleTestWriter creates an option that correctly sets the file frame depth for testing.TB log.
+func ConsoleTestWriter(t testing.TB) func(w *ConsoleWriter) {
+	return func(w *ConsoleWriter) {
+		w.Out = TestWriter{TB: t, Frame: 7}
+	}
+}

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"
@@ -53,45 +55,45 @@ func TestResilientMultiWriter(t *testing.T) {
 		writers []io.Writer
 	}{
 		{
-			name:    "All valid writers",
+			name: "All valid writers",
 			writers: []io.Writer{
-				mockedWriter {
+				mockedWriter{
 					wantErr: false,
 				},
-				mockedWriter {
+				mockedWriter{
 					wantErr: false,
 				},
 			},
 		},
 		{
-			name:    "All invalid writers",
+			name: "All invalid writers",
 			writers: []io.Writer{
-				mockedWriter {
+				mockedWriter{
 					wantErr: true,
 				},
-				mockedWriter {
+				mockedWriter{
 					wantErr: true,
 				},
 			},
 		},
 		{
-			name:    "First invalid writer",
+			name: "First invalid writer",
 			writers: []io.Writer{
-				mockedWriter {
+				mockedWriter{
 					wantErr: true,
 				},
-				mockedWriter {
+				mockedWriter{
 					wantErr: false,
 				},
 			},
 		},
 		{
-			name:    "First valid writer",
+			name: "First valid writer",
 			writers: []io.Writer{
-				mockedWriter {
+				mockedWriter{
 					wantErr: false,
 				},
-				mockedWriter {
+				mockedWriter{
 					wantErr: true,
 				},
 			},
@@ -110,4 +112,66 @@ func TestResilientMultiWriter(t *testing.T) {
 		}
 		writeCalls = 0
 	}
-}
+}
+
+type tb struct {
+	testing.TB
+	buf bytes.Buffer
+}
+
+func (t *tb) Log(args ...interface{}) {
+	if _, err := t.buf.WriteString(fmt.Sprintln(args...)); err != nil {
+		t.Error(err)
+	}
+}
+
+func (t *tb) 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 := &tb{TB: t} // Capture TB log buffer.
+			w := TestWriter{TB: 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()
+		})
+	}
+
+}