chore: simplify tests with logtest (#3184)

core/logc/logs_test.go

@@ -1,7 +1,6 @@
 package logc
 
 import (
-	"bytes"
 	"context"
 	"encoding/json"
 	"fmt"
@@ -11,14 +10,11 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/zeromicro/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/logx/logtest"
 )
 
 func TestAddGlobalFields(t *testing.T) {
-	var buf bytes.Buffer
-	writer := logx.NewWriter(&buf)
-	old := logx.Reset()
-	logx.SetWriter(writer)
-	defer logx.SetWriter(old)
+	buf := logtest.NewCollector(t)
 
 	Info(context.Background(), "hello")
 	buf.Reset()
@@ -34,155 +30,90 @@ func TestAddGlobalFields(t *testing.T) {
 }
 
 func TestAlert(t *testing.T) {
-	var buf strings.Builder
-	writer := logx.NewWriter(&buf)
-	old := logx.Reset()
-	logx.SetWriter(writer)
-	defer logx.SetWriter(old)
-
+	buf := logtest.NewCollector(t)
 	Alert(context.Background(), "foo")
 	assert.True(t, strings.Contains(buf.String(), "foo"), buf.String())
 }
 
 func TestError(t *testing.T) {
-	var buf strings.Builder
-	writer := logx.NewWriter(&buf)
-	old := logx.Reset()
-	logx.SetWriter(writer)
-	defer logx.SetWriter(old)
-
+	buf := logtest.NewCollector(t)
 	file, line := getFileLine()
 	Error(context.Background(), "foo")
 	assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)))
 }
 
 func TestErrorf(t *testing.T) {
-	var buf strings.Builder
-	writer := logx.NewWriter(&buf)
-	old := logx.Reset()
-	logx.SetWriter(writer)
-	defer logx.SetWriter(old)
-
+	buf := logtest.NewCollector(t)
 	file, line := getFileLine()
 	Errorf(context.Background(), "foo %s", "bar")
 	assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)))
 }
 
 func TestErrorv(t *testing.T) {
-	var buf strings.Builder
-	writer := logx.NewWriter(&buf)
-	old := logx.Reset()
-	logx.SetWriter(writer)
-	defer logx.SetWriter(old)
-
+	buf := logtest.NewCollector(t)
 	file, line := getFileLine()
 	Errorv(context.Background(), "foo")
 	assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)))
 }
 
 func TestErrorw(t *testing.T) {
-	var buf strings.Builder
-	writer := logx.NewWriter(&buf)
-	old := logx.Reset()
-	logx.SetWriter(writer)
-	defer logx.SetWriter(old)
-
+	buf := logtest.NewCollector(t)
 	file, line := getFileLine()
 	Errorw(context.Background(), "foo", Field("a", "b"))
 	assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)))
 }
 
 func TestInfo(t *testing.T) {
-	var buf strings.Builder
-	writer := logx.NewWriter(&buf)
-	old := logx.Reset()
-	logx.SetWriter(writer)
-	defer logx.SetWriter(old)
-
+	buf := logtest.NewCollector(t)
 	file, line := getFileLine()
 	Info(context.Background(), "foo")
 	assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)))
 }
 
 func TestInfof(t *testing.T) {
-	var buf strings.Builder
-	writer := logx.NewWriter(&buf)
-	old := logx.Reset()
-	logx.SetWriter(writer)
-	defer logx.SetWriter(old)
-
+	buf := logtest.NewCollector(t)
 	file, line := getFileLine()
 	Infof(context.Background(), "foo %s", "bar")
 	assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)))
 }
 
 func TestInfov(t *testing.T) {
-	var buf strings.Builder
-	writer := logx.NewWriter(&buf)
-	old := logx.Reset()
-	logx.SetWriter(writer)
-	defer logx.SetWriter(old)
-
+	buf := logtest.NewCollector(t)
 	file, line := getFileLine()
 	Infov(context.Background(), "foo")
 	assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)))
 }
 
 func TestInfow(t *testing.T) {
-	var buf strings.Builder
-	writer := logx.NewWriter(&buf)
-	old := logx.Reset()
-	logx.SetWriter(writer)
-	defer logx.SetWriter(old)
-
+	buf := logtest.NewCollector(t)
 	file, line := getFileLine()
 	Infow(context.Background(), "foo", Field("a", "b"))
 	assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)))
 }
 
 func TestDebug(t *testing.T) {
-	var buf strings.Builder
-	writer := logx.NewWriter(&buf)
-	old := logx.Reset()
-	logx.SetWriter(writer)
-	defer logx.SetWriter(old)
-
+	buf := logtest.NewCollector(t)
 	file, line := getFileLine()
 	Debug(context.Background(), "foo")
 	assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)))
 }
 
 func TestDebugf(t *testing.T) {
-	var buf strings.Builder
-	writer := logx.NewWriter(&buf)
-	old := logx.Reset()
-	logx.SetWriter(writer)
-	defer logx.SetWriter(old)
-
+	buf := logtest.NewCollector(t)
 	file, line := getFileLine()
 	Debugf(context.Background(), "foo %s", "bar")
 	assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)))
 }
 
 func TestDebugv(t *testing.T) {
-	var buf strings.Builder
-	writer := logx.NewWriter(&buf)
-	old := logx.Reset()
-	logx.SetWriter(writer)
-	defer logx.SetWriter(old)
-
+	buf := logtest.NewCollector(t)
 	file, line := getFileLine()
 	Debugv(context.Background(), "foo")
 	assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)))
 }
 
 func TestDebugw(t *testing.T) {
-	var buf strings.Builder
-	writer := logx.NewWriter(&buf)
-	old := logx.Reset()
-	logx.SetWriter(writer)
-	defer logx.SetWriter(old)
-
+	buf := logtest.NewCollector(t)
 	file, line := getFileLine()
 	Debugw(context.Background(), "foo", Field("a", "b"))
 	assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)))
@@ -204,48 +135,28 @@ func TestMisc(t *testing.T) {
 }
 
 func TestSlow(t *testing.T) {
-	var buf strings.Builder
-	writer := logx.NewWriter(&buf)
-	old := logx.Reset()
-	logx.SetWriter(writer)
-	defer logx.SetWriter(old)
-
+	buf := logtest.NewCollector(t)
 	file, line := getFileLine()
 	Slow(context.Background(), "foo")
 	assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)), buf.String())
 }
 
 func TestSlowf(t *testing.T) {
-	var buf strings.Builder
-	writer := logx.NewWriter(&buf)
-	old := logx.Reset()
-	logx.SetWriter(writer)
-	defer logx.SetWriter(old)
-
+	buf := logtest.NewCollector(t)
 	file, line := getFileLine()
 	Slowf(context.Background(), "foo %s", "bar")
 	assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)), buf.String())
 }
 
 func TestSlowv(t *testing.T) {
-	var buf strings.Builder
-	writer := logx.NewWriter(&buf)
-	old := logx.Reset()
-	logx.SetWriter(writer)
-	defer logx.SetWriter(old)
-
+	buf := logtest.NewCollector(t)
 	file, line := getFileLine()
 	Slowv(context.Background(), "foo")
 	assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)), buf.String())
 }
 
 func TestSloww(t *testing.T) {
-	var buf strings.Builder
-	writer := logx.NewWriter(&buf)
-	old := logx.Reset()
-	logx.SetWriter(writer)
-	defer logx.SetWriter(old)
-
+	buf := logtest.NewCollector(t)
 	file, line := getFileLine()
 	Sloww(context.Background(), "foo", Field("a", "b"))
 	assert.True(t, strings.Contains(buf.String(), fmt.Sprintf("%s:%d", file, line+1)), buf.String())

core/logx/logtest/logtest.go

@@ -0,0 +1,78 @@
+package logtest
+
+import (
+	"bytes"
+	"encoding/json"
+	"io"
+	"testing"
+
+	"github.com/zeromicro/go-zero/core/logx"
+)
+
+type Buffer struct {
+	buf *bytes.Buffer
+	t   *testing.T
+}
+
+func Discard(t *testing.T) {
+	prev := logx.Reset()
+	logx.SetWriter(logx.NewWriter(io.Discard))
+
+	t.Cleanup(func() {
+		logx.SetWriter(prev)
+	})
+}
+
+func NewCollector(t *testing.T) *Buffer {
+	var buf bytes.Buffer
+	writer := logx.NewWriter(&buf)
+	prev := logx.Reset()
+	logx.SetWriter(writer)
+
+	t.Cleanup(func() {
+		logx.SetWriter(prev)
+	})
+
+	return &Buffer{
+		buf: &buf,
+		t:   t,
+	}
+}
+
+func (b *Buffer) Bytes() []byte {
+	return b.buf.Bytes()
+}
+
+func (b *Buffer) Content() string {
+	var m map[string]interface{}
+	if err := json.Unmarshal(b.buf.Bytes(), &m); err != nil {
+		b.t.Error(err)
+		return ""
+	}
+
+	content, ok := m["content"]
+	if !ok {
+		return ""
+	}
+
+	switch val := content.(type) {
+	case string:
+		return val
+	default:
+		bs, err := json.Marshal(content)
+		if err != nil {
+			b.t.Error(err)
+			return ""
+		}
+
+		return string(bs)
+	}
+}
+
+func (b *Buffer) Reset() {
+	b.buf.Reset()
+}
+
+func (b *Buffer) String() string {
+	return b.buf.String()
+}

core/logx/logtest/logtest_test.go

@@ -0,0 +1,22 @@
+package logtest
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/zeromicro/go-zero/core/logx"
+)
+
+func TestCollector(t *testing.T) {
+	const input = "hello"
+	c := NewCollector(t)
+	logx.Info(input)
+	assert.Equal(t, input, c.Content())
+	assert.Contains(t, c.String(), input)
+}
+
+func TestDiscard(t *testing.T) {
+	const input = "hello"
+	Discard(t)
+	logx.Info(input)
+}

core/proc/goroutines_test.go

@@ -5,19 +5,11 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/zeromicro/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/logx/logtest"
 )
 
 func TestDumpGoroutines(t *testing.T) {
-	var buf strings.Builder
-	w := logx.NewWriter(&buf)
-	o := logx.Reset()
-	logx.SetWriter(w)
-	defer func() {
-		logx.Reset()
-		logx.SetWriter(o)
-	}()
-
+	buf := logtest.NewCollector(t)
 	dumpGoroutines()
 	assert.True(t, strings.Contains(buf.String(), ".dump"))
 }

core/proc/profile_test.go

@@ -5,25 +5,16 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/zeromicro/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/logx/logtest"
 )
 
 func TestProfile(t *testing.T) {
-	var buf strings.Builder
-	w := logx.NewWriter(&buf)
-	o := logx.Reset()
-	logx.SetWriter(w)
-
-	defer func() {
-		logx.Reset()
-		logx.SetWriter(o)
-	}()
-
+	c := logtest.NewCollector(t)
 	profiler := StartProfile()
 	// start again should not work
 	assert.NotNil(t, StartProfile())
 	profiler.Stop()
 	// stop twice
 	profiler.Stop()
-	assert.True(t, strings.Contains(buf.String(), ".pprof"))
+	assert.True(t, strings.Contains(c.String(), ".pprof"))
 }

core/stat/usage_test.go

@@ -1,12 +1,11 @@
 package stat
 
 import (
-	"bytes"
 	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/zeromicro/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/logx/logtest"
 )
 
 func TestBToMb(t *testing.T) {
@@ -41,15 +40,11 @@ func TestBToMb(t *testing.T) {
 }
 
 func TestPrintUsage(t *testing.T) {
-	var buf bytes.Buffer
-	writer := logx.NewWriter(&buf)
-	old := logx.Reset()
-	logx.SetWriter(writer)
-	defer logx.SetWriter(old)
+	c := logtest.NewCollector(t)
 
 	printUsage()
 
-	output := buf.String()
+	output := c.String()
 	assert.Contains(t, output, "CPU:")
 	assert.Contains(t, output, "MEMORY:")
 	assert.Contains(t, output, "Alloc=")

core/stores/mon/collection_test.go

@@ -3,12 +3,11 @@ package mon
 import (
 	"context"
 	"errors"
-	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/zeromicro/go-zero/core/breaker"
-	"github.com/zeromicro/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/logx/logtest"
 	"github.com/zeromicro/go-zero/core/stringx"
 	"github.com/zeromicro/go-zero/core/timex"
 	"go.mongodb.org/mongo-driver/bson"
@@ -573,15 +572,7 @@ func TestDecoratedCollection_LogDuration(t *testing.T) {
 		brk:        breaker.NewBreaker(),
 	}
 
-	var buf strings.Builder
-	w := logx.NewWriter(&buf)
-	o := logx.Reset()
-	logx.SetWriter(w)
-
-	defer func() {
-		logx.Reset()
-		logx.SetWriter(o)
-	}()
+	buf := logtest.NewCollector(t)
 
 	buf.Reset()
 	c.logDuration(context.Background(), "foo", timex.Now(), nil, "bar")

core/stores/mon/util_test.go

@@ -3,12 +3,11 @@ package mon
 import (
 	"context"
 	"errors"
-	"strings"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/zeromicro/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/logx/logtest"
 )
 
 func TestFormatAddrs(t *testing.T) {
@@ -40,15 +39,7 @@ func TestFormatAddrs(t *testing.T) {
 }
 
 func Test_logDuration(t *testing.T) {
-	var buf strings.Builder
-	w := logx.NewWriter(&buf)
-	o := logx.Reset()
-	logx.SetWriter(w)
-
-	defer func() {
-		logx.Reset()
-		logx.SetWriter(o)
-	}()
+	buf := logtest.NewCollector(t)
 
 	buf.Reset()
 	logDuration(context.Background(), "foo", "bar", time.Millisecond, nil)

core/stores/redis/hook_test.go

@@ -9,7 +9,7 @@ import (
 
 	red "github.com/go-redis/redis/v8"
 	"github.com/stretchr/testify/assert"
-	"github.com/zeromicro/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/logx/logtest"
 	ztrace "github.com/zeromicro/go-zero/core/trace"
 	tracesdk "go.opentelemetry.io/otel/trace"
 )
@@ -47,8 +47,7 @@ func TestHookProcessCase2(t *testing.T) {
 	})
 	defer ztrace.StopAgent()
 
-	w, restore := injectLog()
-	defer restore()
+	w := logtest.NewCollector(t)
 
 	ctx, err := durationHook.BeforeProcess(context.Background(), red.NewCmd(context.Background()))
 	if err != nil {
@@ -115,8 +114,7 @@ func TestHookProcessPipelineCase2(t *testing.T) {
 	})
 	defer ztrace.StopAgent()
 
-	w, restore := injectLog()
-	defer restore()
+	w := logtest.NewCollector(t)
 
 	ctx, err := durationHook.BeforeProcessPipeline(context.Background(), []red.Cmder{
 		red.NewCmd(context.Background()),
@@ -135,8 +133,7 @@ func TestHookProcessPipelineCase2(t *testing.T) {
 }
 
 func TestHookProcessPipelineCase3(t *testing.T) {
-	w, restore := injectLog()
-	defer restore()
+	w := logtest.NewCollector(t)
 
 	assert.Nil(t, durationHook.AfterProcessPipeline(context.Background(), []red.Cmder{
 		red.NewCmd(context.Background()),
@@ -145,8 +142,7 @@ func TestHookProcessPipelineCase3(t *testing.T) {
 }
 
 func TestHookProcessPipelineCase4(t *testing.T) {
-	w, restore := injectLog()
-	defer restore()
+	w := logtest.NewCollector(t)
 
 	ctx := context.WithValue(context.Background(), startTimeKey, "foo")
 	assert.Nil(t, durationHook.AfterProcessPipeline(ctx, []red.Cmder{
@@ -169,8 +165,7 @@ func TestHookProcessPipelineCase5(t *testing.T) {
 }
 
 func TestLogDuration(t *testing.T) {
-	w, restore := injectLog()
-	defer restore()
+	w := logtest.NewCollector(t)
 
 	logDuration(context.Background(), []red.Cmder{
 		red.NewCmd(context.Background(), "get", "foo"),
@@ -183,15 +178,3 @@ func TestLogDuration(t *testing.T) {
 	}, 1*time.Second)
 	assert.True(t, strings.Contains(w.String(), `get foo\nset bar 0`))
 }
-
-func injectLog() (r *strings.Builder, restore func()) {
-	var buf strings.Builder
-	w := logx.NewWriter(&buf)
-	o := logx.Reset()
-	logx.SetWriter(w)
-
-	return &buf, func() {
-		logx.Reset()
-		logx.SetWriter(o)
-	}
-}

core/trace/tracetest/tracetest.go

@@ -16,5 +16,6 @@ func NewInMemoryExporter(t *testing.T) *tracetest.InMemoryExporter {
 		me.Reset()
 	})
 	otel.SetTracerProvider(trace.NewTracerProvider(trace.WithSyncer(me)))
+
 	return me
 }