chore: simplify tests with logtest (#3184)

zrpc/internal/rpclogger_test.go

@@ -1,121 +1,96 @@
 package internal
 
 import (
-	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/zeromicro/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/logx/logtest"
 )
 
 const content = "foo"
 
 func TestLoggerError(t *testing.T) {
-	w, restore := injectLog()
-	defer restore()
-
+	c := logtest.NewCollector(t)
 	logger := new(Logger)
 	logger.Error(content)
-	assert.Contains(t, w.String(), content)
+	assert.Contains(t, c.String(), content)
 }
 
 func TestLoggerErrorf(t *testing.T) {
-	w, restore := injectLog()
-	defer restore()
-
+	c := logtest.NewCollector(t)
 	logger := new(Logger)
 	logger.Errorf(content)
-	assert.Contains(t, w.String(), content)
+	assert.Contains(t, c.String(), content)
 }
 
 func TestLoggerErrorln(t *testing.T) {
-	w, restore := injectLog()
-	defer restore()
-
+	c := logtest.NewCollector(t)
 	logger := new(Logger)
 	logger.Errorln(content)
-	assert.Contains(t, w.String(), content)
+	assert.Contains(t, c.String(), content)
 }
 
 func TestLoggerFatal(t *testing.T) {
-	w, restore := injectLog()
-	defer restore()
-
+	c := logtest.NewCollector(t)
 	logger := new(Logger)
 	logger.Fatal(content)
-	assert.Contains(t, w.String(), content)
+	assert.Contains(t, c.String(), content)
 }
 
 func TestLoggerFatalf(t *testing.T) {
-	w, restore := injectLog()
-	defer restore()
-
+	c := logtest.NewCollector(t)
 	logger := new(Logger)
 	logger.Fatalf(content)
-	assert.Contains(t, w.String(), content)
+	assert.Contains(t, c.String(), content)
 }
 
 func TestLoggerFatalln(t *testing.T) {
-	w, restore := injectLog()
-	defer restore()
-
+	c := logtest.NewCollector(t)
 	logger := new(Logger)
 	logger.Fatalln(content)
-	assert.Contains(t, w.String(), content)
+	assert.Contains(t, c.String(), content)
 }
 
 func TestLoggerInfo(t *testing.T) {
-	w, restore := injectLog()
-	defer restore()
-
+	c := logtest.NewCollector(t)
 	logger := new(Logger)
 	logger.Info(content)
-	assert.Empty(t, w.String())
+	assert.Empty(t, c.String())
 }
 
 func TestLoggerInfof(t *testing.T) {
-	w, restore := injectLog()
-	defer restore()
-
+	c := logtest.NewCollector(t)
 	logger := new(Logger)
 	logger.Infof(content)
-	assert.Empty(t, w.String())
+	assert.Empty(t, c.String())
 }
 
 func TestLoggerWarning(t *testing.T) {
-	w, restore := injectLog()
-	defer restore()
-
+	c := logtest.NewCollector(t)
 	logger := new(Logger)
 	logger.Warning(content)
-	assert.Empty(t, w.String())
+	assert.Empty(t, c.String())
 }
 
 func TestLoggerInfoln(t *testing.T) {
-	w, restore := injectLog()
-	defer restore()
-
+	c := logtest.NewCollector(t)
 	logger := new(Logger)
 	logger.Infoln(content)
-	assert.Empty(t, w.String())
+	assert.Empty(t, c.String())
 }
 
 func TestLoggerWarningf(t *testing.T) {
-	w, restore := injectLog()
-	defer restore()
-
+	c := logtest.NewCollector(t)
 	logger := new(Logger)
 	logger.Warningf(content)
-	assert.Empty(t, w.String())
+	assert.Empty(t, c.String())
 }
 
 func TestLoggerWarningln(t *testing.T) {
-	w, restore := injectLog()
-	defer restore()
-
+	c := logtest.NewCollector(t)
 	logger := new(Logger)
 	logger.Warningln(content)
-	assert.Empty(t, w.String())
+	assert.Empty(t, c.String())
 }
 
 func TestLogger_V(t *testing.T) {
@@ -125,15 +100,3 @@ func TestLogger_V(t *testing.T) {
 	// grpclog.infoLog
 	assert.False(t, logger.V(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)
-	}
-}

zrpc/internal/rpcserver_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/zeromicro/go-zero/core/proc"
@@ -40,6 +41,8 @@ func TestRpcServer(t *testing.T) {
 	}()
 
 	wg.Wait()
+	time.Sleep(100 * time.Millisecond)
+
 	lock.Lock()
 	grpcServer.GracefulStop()
 	lock.Unlock()