From 4c68b7df7cdea1c4404a53d87852b7c4ce9bd158 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Tue, 7 Sep 2021 15:17:39 -0700
Subject: [PATCH] tstest: add MemLogger bytes.Buffer wrapper with Logf method

We use it tons of places. Updated three at least in this PR.

Another use in next commit.

Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 tstest/log.go                        | 14 ++++++++++++++
 wgengine/router/router_linux_test.go | 11 +++--------
 wgengine/userspace_test.go           | 12 +++++-------
 wgengine/watchdog_test.go            |  9 ++++-----
 4 files changed, 26 insertions(+), 20 deletions(-)

diff --git a/tstest/log.go b/tstest/log.go
index 8443aade1..867a5a1bc 100644
--- a/tstest/log.go
+++ b/tstest/log.go
@@ -12,6 +12,7 @@ import (
 	"sync"
 	"testing"
 
+	"go4.org/mem"
 	"tailscale.com/types/logger"
 )
 
@@ -122,3 +123,16 @@ func (lt *LogLineTracker) Close() {
 	defer lt.mu.Unlock()
 	lt.closed = true
 }
+
+// MemLogger is a bytes.Buffer with a Logf method for tests that want
+// to log to a buffer.
+type MemLogger struct {
+	bytes.Buffer
+}
+
+func (ml *MemLogger) Logf(format string, args ...interface{}) {
+	fmt.Fprintf(&ml.Buffer, format, args...)
+	if !mem.HasSuffix(mem.B(ml.Buffer.Bytes()), mem.S("\n")) {
+		ml.Buffer.WriteByte('\n')
+	}
+}
diff --git a/wgengine/router/router_linux_test.go b/wgengine/router/router_linux_test.go
index 8b9b228bb..ab42752c8 100644
--- a/wgengine/router/router_linux_test.go
+++ b/wgengine/router/router_linux_test.go
@@ -5,7 +5,6 @@
 package router
 
 import (
-	"bytes"
 	"errors"
 	"fmt"
 	"math/rand"
@@ -18,6 +17,7 @@ import (
 	"github.com/google/go-cmp/cmp"
 	"golang.zx2c4.com/wireguard/tun"
 	"inet.af/netaddr"
+	"tailscale.com/tstest"
 	"tailscale.com/types/logger"
 	"tailscale.com/wgengine/monitor"
 )
@@ -660,13 +660,8 @@ func TestDelRouteIdempotent(t *testing.T) {
 	tun := createTestTUN(t)
 	defer tun.Close()
 
-	var logOutput bytes.Buffer
-	logf := func(format string, args ...interface{}) {
-		fmt.Fprintf(&logOutput, format, args...)
-		if !bytes.HasSuffix(logOutput.Bytes(), []byte("\n")) {
-			logOutput.WriteByte('\n')
-		}
-	}
+	var logOutput tstest.MemLogger
+	logf := logOutput.Logf
 
 	mon, err := monitor.New(logger.Discard)
 	if err != nil {
diff --git a/wgengine/userspace_test.go b/wgengine/userspace_test.go
index e6e2b4115..bb4bd7efb 100644
--- a/wgengine/userspace_test.go
+++ b/wgengine/userspace_test.go
@@ -5,7 +5,6 @@
 package wgengine
 
 import (
-	"bytes"
 	"fmt"
 	"reflect"
 	"testing"
@@ -15,6 +14,7 @@ import (
 	"tailscale.com/net/dns"
 	"tailscale.com/net/tstun"
 	"tailscale.com/tailcfg"
+	"tailscale.com/tstest"
 	"tailscale.com/tstime/mono"
 	"tailscale.com/types/key"
 	"tailscale.com/types/netmap"
@@ -25,7 +25,7 @@ import (
 
 func TestNoteReceiveActivity(t *testing.T) {
 	now := mono.Time(123456)
-	var logBuf bytes.Buffer
+	var logBuf tstest.MemLogger
 
 	confc := make(chan bool, 1)
 	gotConf := func() bool {
@@ -37,11 +37,9 @@ func TestNoteReceiveActivity(t *testing.T) {
 		}
 	}
 	e := &userspaceEngine{
-		timeNow:        func() mono.Time { return now },
-		recvActivityAt: map[tailcfg.NodeKey]mono.Time{},
-		logf: func(format string, a ...interface{}) {
-			fmt.Fprintf(&logBuf, format, a...)
-		},
+		timeNow:               func() mono.Time { return now },
+		recvActivityAt:        map[tailcfg.NodeKey]mono.Time{},
+		logf:                  logBuf.Logf,
 		tundev:                new(tstun.Wrapper),
 		testMaybeReconfigHook: func() { confc <- true },
 		trimmedNodes:          map[tailcfg.NodeKey]bool{},
diff --git a/wgengine/watchdog_test.go b/wgengine/watchdog_test.go
index 0a699f848..05ca79d7d 100644
--- a/wgengine/watchdog_test.go
+++ b/wgengine/watchdog_test.go
@@ -5,12 +5,13 @@
 package wgengine
 
 import (
-	"bytes"
 	"fmt"
 	"runtime"
 	"strings"
 	"testing"
 	"time"
+
+	"tailscale.com/tstest"
 )
 
 func TestWatchdog(t *testing.T) {
@@ -52,11 +53,9 @@ func TestWatchdog(t *testing.T) {
 		wdEngine := e.(*watchdogEngine)
 		wdEngine.maxWait = maxWaitMultiple * 100 * time.Millisecond
 
-		logBuf := new(bytes.Buffer)
+		logBuf := new(tstest.MemLogger)
 		fatalCalled := make(chan struct{})
-		wdEngine.logf = func(format string, args ...interface{}) {
-			fmt.Fprintf(logBuf, format+"\n", args...)
-		}
+		wdEngine.logf = logBuf.Logf
 		wdEngine.fatalf = func(format string, args ...interface{}) {
 			t.Logf("FATAL: %s", fmt.Sprintf(format, args...))
 			fatalCalled <- struct{}{}