avoid bigint converted into float64 when unmarshaling

* rebase upstream

* rebase

* trim no need line

* trim no need line

* trim no need line

* update doc

* remove update

* remove no need

* remove no need

* goctl add jwt support

* goctl add jwt support

* goctl add jwt support

* goctl support import

* goctl support import

* support return ()

* revert

* refactor and rename folder to group

* parser add test

Co-authored-by: kingxt <dream4kingxt@163.com>

core/cmdline/input_test.go

@@ -8,20 +8,14 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/iox"
 	"github.com/tal-tech/go-zero/core/lang"
 )
 
 func TestEnterToContinue(t *testing.T) {
-	r, w, err := os.Pipe()
+	restore, err := iox.RedirectInOut()
 	assert.Nil(t, err)
-	ow := os.Stdout
-	os.Stdout = w
-	or := os.Stdin
-	os.Stdin = r
-	defer func() {
-		os.Stdin = or
-		os.Stdout = ow
-	}()
+	defer restore()
 
 	var wg sync.WaitGroup
 	wg.Add(2)

core/discov/internal/listener.go

@@ -1,4 +1,3 @@
-//go:generate mockgen -package internal -destination listener_mock.go -source listener.go Listener
 package internal
 
 type Listener interface {

core/discov/internal/listener_mock.go

@@ -1,45 +0,0 @@
-// Code generated by MockGen. DO NOT EDIT.
-// Source: listener.go
-
-// Package internal is a generated GoMock package.
-package internal
-
-import (
-	gomock "github.com/golang/mock/gomock"
-	reflect "reflect"
-)
-
-// MockListener is a mock of Listener interface
-type MockListener struct {
-	ctrl     *gomock.Controller
-	recorder *MockListenerMockRecorder
-}
-
-// MockListenerMockRecorder is the mock recorder for MockListener
-type MockListenerMockRecorder struct {
-	mock *MockListener
-}
-
-// NewMockListener creates a new mock instance
-func NewMockListener(ctrl *gomock.Controller) *MockListener {
-	mock := &MockListener{ctrl: ctrl}
-	mock.recorder = &MockListenerMockRecorder{mock}
-	return mock
-}
-
-// EXPECT returns an object that allows the caller to indicate expected use
-func (m *MockListener) EXPECT() *MockListenerMockRecorder {
-	return m.recorder
-}
-
-// OnUpdate mocks base method
-func (m *MockListener) OnUpdate(keys, values []string, newKey string) {
-	m.ctrl.T.Helper()
-	m.ctrl.Call(m, "OnUpdate", keys, values, newKey)
-}
-
-// OnUpdate indicates an expected call of OnUpdate
-func (mr *MockListenerMockRecorder) OnUpdate(keys, values, newKey interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnUpdate", reflect.TypeOf((*MockListener)(nil).OnUpdate), keys, values, newKey)
-}

core/discov/publisher_test.go

@@ -111,6 +111,10 @@ func TestPublisher_keepAliveAsyncQuit(t *testing.T) {
 	defer ctrl.Finish()
 	const id clientv3.LeaseID = 1
 	cli := internal.NewMockEtcdClient(ctrl)
+	cli.EXPECT().ActiveConnection()
+	cli.EXPECT().Close()
+	defer cli.Close()
+	cli.ActiveConnection()
 	restore := setMockClient(cli)
 	defer restore()
 	cli.EXPECT().Ctx().AnyTimes()

core/fx/fn.go

@@ -84,6 +84,14 @@ func (p Stream) Buffer(n int) Stream {
 	return Range(source)
 }
 
+// Count counts the number of elements in the result.
+func (p Stream) Count() (count int) {
+	for range p.source {
+		count++
+	}
+	return
+}
+
 // Distinct removes the duplicated items base on the given KeyFunc.
 func (p Stream) Distinct(fn KeyFunc) Stream {
 	source := make(chan interface{})

core/fx/fn_test.go

@@ -49,6 +49,36 @@ func TestBufferNegative(t *testing.T) {
 	assert.Equal(t, 10, result)
 }
 
+func TestCount(t *testing.T) {
+	tests := []struct {
+		name     string
+		elements []interface{}
+	}{
+		{
+			name: "no elements with nil",
+		},
+		{
+			name:     "no elements",
+			elements: []interface{}{},
+		},
+		{
+			name:     "1 element",
+			elements: []interface{}{1},
+		},
+		{
+			name:     "multiple elements",
+			elements: []interface{}{1, 2, 3},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			val := Just(test.elements...).Count()
+			assert.Equal(t, len(test.elements), val)
+		})
+	}
+}
+
 func TestDone(t *testing.T) {
 	var count int32
 	Just(1, 2, 3).Walk(func(item interface{}, pipe chan<- interface{}) {

core/iox/pipe.go

@@ -0,0 +1,23 @@
+package iox
+
+import "os"
+
+// RedirectInOut redirects stdin to r, stdout to w, and callers need to call restore afterwards.
+func RedirectInOut() (restore func(), err error) {
+	var r, w *os.File
+	r, w, err = os.Pipe()
+	if err != nil {
+		return
+	}
+
+	ow := os.Stdout
+	os.Stdout = w
+	or := os.Stdin
+	os.Stdin = r
+	restore = func() {
+		os.Stdin = or
+		os.Stdout = ow
+	}
+
+	return
+}

core/iox/pipe_test.go

@@ -0,0 +1,13 @@
+package iox
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRedirectInOut(t *testing.T) {
+	restore, err := RedirectInOut()
+	assert.Nil(t, err)
+	defer restore()
+}

core/load/adaptiveshedder_test.go

@@ -135,6 +135,7 @@ func TestAdaptiveShedderShouldDrop(t *testing.T) {
 		passCounter:     passCounter,
 		rtCounter:       rtCounter,
 		windows:         buckets,
+		dropTime:        syncx.NewAtomicDuration(),
 		droppedRecently: syncx.NewAtomicBool(),
 	}
 	// cpu >=  800, inflight < maxPass
@@ -160,6 +161,40 @@ func TestAdaptiveShedderShouldDrop(t *testing.T) {
 	}
 	shedder.avgFlying = 80
 	assert.False(t, shedder.shouldDrop())
+
+	// cpu >=  800, inflight < maxPass
+	systemOverloadChecker = func(int64) bool {
+		return true
+	}
+	shedder.avgFlying = 80
+	shedder.flying = 80
+	_, err := shedder.Allow()
+	assert.NotNil(t, err)
+}
+
+func TestAdaptiveShedderStillHot(t *testing.T) {
+	logx.Disable()
+	passCounter := newRollingWindow()
+	rtCounter := newRollingWindow()
+	for i := 0; i < 10; i++ {
+		if i > 0 {
+			time.Sleep(bucketDuration)
+		}
+		passCounter.Add(float64((i + 1) * 100))
+		for j := i*10 + 1; j <= i*10+10; j++ {
+			rtCounter.Add(float64(j))
+		}
+	}
+	shedder := &adaptiveShedder{
+		passCounter:     passCounter,
+		rtCounter:       rtCounter,
+		windows:         buckets,
+		dropTime:        syncx.NewAtomicDuration(),
+		droppedRecently: syncx.ForAtomicBool(true),
+	}
+	assert.False(t, shedder.stillHot())
+	shedder.dropTime.Set(-coolOffDuration * 2)
+	assert.False(t, shedder.stillHot())
 }
 
 func BenchmarkAdaptiveShedder_Allow(b *testing.B) {

core/load/sheddergroup_test.go

@@ -13,3 +13,8 @@ func TestGroup(t *testing.T) {
 		assert.NotNil(t, limiter)
 	})
 }
+
+func TestShedderClose(t *testing.T) {
+	var nop nopCloser
+	assert.Nil(t, nop.Close())
+}

core/logx/logs_test.go

@@ -6,8 +6,10 @@ import (
 	"io"
 	"io/ioutil"
 	"log"
+	"os"
 	"runtime"
 	"strings"
+	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -21,10 +23,13 @@ var (
 )
 
 type mockWriter struct {
+	lock    sync.Mutex
 	builder strings.Builder
 }
 
 func (mw *mockWriter) Write(data []byte) (int, error) {
+	mw.lock.Lock()
+	defer mw.lock.Unlock()
 	return mw.builder.Write(data)
 }
 
@@ -32,12 +37,22 @@ func (mw *mockWriter) Close() error {
 	return nil
 }
 
+func (mw *mockWriter) Contains(text string) bool {
+	mw.lock.Lock()
+	defer mw.lock.Unlock()
+	return strings.Contains(mw.builder.String(), text)
+}
+
 func (mw *mockWriter) Reset() {
+	mw.lock.Lock()
+	defer mw.lock.Unlock()
 	mw.builder.Reset()
 }
 
-func (mw *mockWriter) Contains(text string) bool {
-	return strings.Contains(mw.builder.String(), text)
+func (mw *mockWriter) String() string {
+	mw.lock.Lock()
+	defer mw.lock.Unlock()
+	return mw.builder.String()
 }
 
 func TestFileLineFileMode(t *testing.T) {
@@ -175,6 +190,43 @@ func TestMustNil(t *testing.T) {
 	Must(nil)
 }
 
+func TestSetup(t *testing.T) {
+	MustSetup(LogConf{
+		ServiceName: "any",
+		Mode:        "console",
+	})
+	MustSetup(LogConf{
+		ServiceName: "any",
+		Mode:        "file",
+		Path:        os.TempDir(),
+	})
+	MustSetup(LogConf{
+		ServiceName: "any",
+		Mode:        "volume",
+		Path:        os.TempDir(),
+	})
+	assert.NotNil(t, setupWithVolume(LogConf{}))
+	assert.NotNil(t, setupWithFiles(LogConf{}))
+	assert.Nil(t, setupWithFiles(LogConf{
+		ServiceName: "any",
+		Path:        os.TempDir(),
+		Compress:    true,
+		KeepDays:    1,
+	}))
+	setupLogLevel(LogConf{
+		Level: levelInfo,
+	})
+	setupLogLevel(LogConf{
+		Level: levelError,
+	})
+	setupLogLevel(LogConf{
+		Level: levelSevere,
+	})
+	_, err := createOutput("")
+	assert.NotNil(t, err)
+	Disable()
+}
+
 func TestDisable(t *testing.T) {
 	Disable()
 	WithKeepDays(1)
@@ -184,6 +236,20 @@ func TestDisable(t *testing.T) {
 	assert.Nil(t, Close())
 }
 
+func TestWithGzip(t *testing.T) {
+	fn := WithGzip()
+	var opt logOptions
+	fn(&opt)
+	assert.True(t, opt.gzipEnabled)
+}
+
+func TestWithKeepDays(t *testing.T) {
+	fn := WithKeepDays(1)
+	var opt logOptions
+	fn(&opt)
+	assert.Equal(t, 1, opt.keepDays)
+}
+
 func BenchmarkCopyByteSliceAppend(b *testing.B) {
 	for i := 0; i < b.N; i++ {
 		var buf []byte
@@ -301,4 +367,10 @@ func testSetLevelTwiceWithMode(t *testing.T, mode string) {
 	atomic.StoreUint32(&initialized, 1)
 	Info(message)
 	assert.Equal(t, 0, writer.builder.Len())
+	Infof(message)
+	assert.Equal(t, 0, writer.builder.Len())
+	ErrorStack(message)
+	assert.Equal(t, 0, writer.builder.Len())
+	ErrorStackf(message)
+	assert.Equal(t, 0, writer.builder.Len())
 }

core/logx/tracelogger_test.go

@@ -2,8 +2,11 @@ package logx
 
 import (
 	"context"
+	"log"
 	"strings"
+	"sync/atomic"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/tal-tech/go-zero/core/trace/tracespec"
@@ -17,13 +20,76 @@ const (
 var mock tracespec.Trace = new(mockTrace)
 
 func TestTraceLog(t *testing.T) {
-	var buf strings.Builder
+	var buf mockWriter
 	ctx := context.WithValue(context.Background(), tracespec.TracingKey, mock)
 	WithContext(ctx).(*traceLogger).write(&buf, levelInfo, testlog)
 	assert.True(t, strings.Contains(buf.String(), mockTraceId))
 	assert.True(t, strings.Contains(buf.String(), mockSpanId))
 }
 
+func TestTraceError(t *testing.T) {
+	var buf mockWriter
+	ctx := context.WithValue(context.Background(), tracespec.TracingKey, mock)
+	l := WithContext(ctx).(*traceLogger)
+	SetLevel(InfoLevel)
+	atomic.StoreUint32(&initialized, 1)
+	errorLog = newLogWriter(log.New(&buf, "", flags))
+	l.WithDuration(time.Second).Error(testlog)
+	assert.True(t, strings.Contains(buf.String(), mockTraceId))
+	assert.True(t, strings.Contains(buf.String(), mockSpanId))
+	buf.Reset()
+	l.WithDuration(time.Second).Errorf(testlog)
+	assert.True(t, strings.Contains(buf.String(), mockTraceId))
+	assert.True(t, strings.Contains(buf.String(), mockSpanId))
+}
+
+func TestTraceInfo(t *testing.T) {
+	var buf mockWriter
+	ctx := context.WithValue(context.Background(), tracespec.TracingKey, mock)
+	l := WithContext(ctx).(*traceLogger)
+	SetLevel(InfoLevel)
+	atomic.StoreUint32(&initialized, 1)
+	infoLog = newLogWriter(log.New(&buf, "", flags))
+	l.WithDuration(time.Second).Info(testlog)
+	assert.True(t, strings.Contains(buf.String(), mockTraceId))
+	assert.True(t, strings.Contains(buf.String(), mockSpanId))
+	buf.Reset()
+	l.WithDuration(time.Second).Infof(testlog)
+	assert.True(t, strings.Contains(buf.String(), mockTraceId))
+	assert.True(t, strings.Contains(buf.String(), mockSpanId))
+}
+
+func TestTraceSlow(t *testing.T) {
+	var buf mockWriter
+	ctx := context.WithValue(context.Background(), tracespec.TracingKey, mock)
+	l := WithContext(ctx).(*traceLogger)
+	SetLevel(InfoLevel)
+	atomic.StoreUint32(&initialized, 1)
+	slowLog = newLogWriter(log.New(&buf, "", flags))
+	l.WithDuration(time.Second).Slow(testlog)
+	assert.True(t, strings.Contains(buf.String(), mockTraceId))
+	assert.True(t, strings.Contains(buf.String(), mockSpanId))
+	buf.Reset()
+	l.WithDuration(time.Second).Slowf(testlog)
+	assert.True(t, strings.Contains(buf.String(), mockTraceId))
+	assert.True(t, strings.Contains(buf.String(), mockSpanId))
+}
+
+func TestTraceWithoutContext(t *testing.T) {
+	var buf mockWriter
+	l := WithContext(context.Background()).(*traceLogger)
+	SetLevel(InfoLevel)
+	atomic.StoreUint32(&initialized, 1)
+	infoLog = newLogWriter(log.New(&buf, "", flags))
+	l.WithDuration(time.Second).Info(testlog)
+	assert.False(t, strings.Contains(buf.String(), mockTraceId))
+	assert.False(t, strings.Contains(buf.String(), mockSpanId))
+	buf.Reset()
+	l.WithDuration(time.Second).Infof(testlog)
+	assert.False(t, strings.Contains(buf.String(), mockTraceId))
+	assert.False(t, strings.Contains(buf.String(), mockSpanId))
+}
+
 type mockTrace struct{}
 
 func (t mockTrace) TraceId() string {

core/stat/alert_test.go

@@ -0,0 +1,22 @@
+// +build linux
+
+package stat
+
+import (
+	"strconv"
+	"sync/atomic"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestReport(t *testing.T) {
+	var count int32
+	SetReporter(func(s string) {
+		atomic.AddInt32(&count, 1)
+	})
+	for i := 0; i < 10; i++ {
+		Report(strconv.Itoa(i))
+	}
+	assert.Equal(t, int32(1), count)
+}

core/stat/internal/cpu_linux_test.go

@@ -1,6 +1,14 @@
 package internal
 
-import "testing"
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRefreshCpu(t *testing.T) {
+	assert.True(t, RefreshCpu() >= 0)
+}
 
 func BenchmarkRefreshCpu(b *testing.B) {
 	for i := 0; i < b.N; i++ {

core/stat/metrics_test.go

@@ -0,0 +1,37 @@
+package stat
+
+import (
+	"strconv"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestMetrics(t *testing.T) {
+	counts := []int{1, 5, 10, 100, 1000, 1000}
+	for _, count := range counts {
+		m := NewMetrics("foo")
+		m.SetName("bar")
+		for i := 0; i < count; i++ {
+			m.Add(Task{
+				Duration:    time.Millisecond * time.Duration(i),
+				Description: strconv.Itoa(i),
+			})
+		}
+		m.AddDrop()
+		var writer mockedWriter
+		SetReportWriter(&writer)
+		m.executor.Flush()
+		assert.Equal(t, "bar", writer.report.Name)
+	}
+}
+
+type mockedWriter struct {
+	report *StatReport
+}
+
+func (m *mockedWriter) Write(report *StatReport) error {
+	m.report = report
+	return nil
+}

core/stat/remotewriter_test.go

@@ -0,0 +1,30 @@
+package stat
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"gopkg.in/h2non/gock.v1"
+)
+
+func TestRemoteWriter(t *testing.T) {
+	defer gock.Off()
+
+	gock.New("http://foo.com").Reply(200).BodyString("foo")
+	writer := NewRemoteWriter("http://foo.com")
+	err := writer.Write(&StatReport{
+		Name: "bar",
+	})
+	assert.Nil(t, err)
+}
+
+func TestRemoteWriterFail(t *testing.T) {
+	defer gock.Off()
+
+	gock.New("http://foo.com").Reply(503).BodyString("foo")
+	writer := NewRemoteWriter("http://foo.com")
+	err := writer.Write(&StatReport{
+		Name: "bar",
+	})
+	assert.NotNil(t, err)
+}

core/stores/cache/cache_test.go

@@ -111,6 +111,45 @@ func TestCache_SetDel(t *testing.T) {
 		assert.Nil(t, c.GetCache(fmt.Sprintf("key/%d", i), &v))
 		assert.Equal(t, i, v)
 	}
+	assert.Nil(t, c.DelCache())
+	for i := 0; i < total; i++ {
+		assert.Nil(t, c.DelCache(fmt.Sprintf("key/%d", i)))
+	}
+	for i := 0; i < total; i++ {
+		var v int
+		assert.Equal(t, errPlaceholder, c.GetCache(fmt.Sprintf("key/%d", i), &v))
+		assert.Equal(t, 0, v)
+	}
+}
+
+func TestCache_OneNode(t *testing.T) {
+	const total = 1000
+	r := miniredis.NewMiniRedis()
+	assert.Nil(t, r.Start())
+	defer r.Close()
+	conf := ClusterConf{
+		{
+			RedisConf: redis.RedisConf{
+				Host: r.Addr(),
+				Type: redis.NodeType,
+			},
+			Weight: 100,
+		},
+	}
+	c := NewCache(conf, syncx.NewSharedCalls(), NewCacheStat("mock"), errPlaceholder)
+	for i := 0; i < total; i++ {
+		if i%2 == 0 {
+			assert.Nil(t, c.SetCache(fmt.Sprintf("key/%d", i), i))
+		} else {
+			assert.Nil(t, c.SetCacheWithExpire(fmt.Sprintf("key/%d", i), i, 0))
+		}
+	}
+	for i := 0; i < total; i++ {
+		var v int
+		assert.Nil(t, c.GetCache(fmt.Sprintf("key/%d", i), &v))
+		assert.Equal(t, i, v)
+	}
+	assert.Nil(t, c.DelCache())
 	for i := 0; i < total; i++ {
 		assert.Nil(t, c.DelCache(fmt.Sprintf("key/%d", i)))
 	}
@@ -188,6 +227,25 @@ func TestCache_Balance(t *testing.T) {
 	assert.Equal(t, total/10, count)
 }
 
+func TestCacheNoNode(t *testing.T) {
+	dispatcher := hash.NewConsistentHash()
+	c := cacheCluster{
+		dispatcher:  dispatcher,
+		errNotFound: errPlaceholder,
+	}
+	assert.NotNil(t, c.DelCache("foo"))
+	assert.NotNil(t, c.DelCache("foo", "bar", "any"))
+	assert.NotNil(t, c.GetCache("foo", nil))
+	assert.NotNil(t, c.SetCache("foo", nil))
+	assert.NotNil(t, c.SetCacheWithExpire("foo", nil, time.Second))
+	assert.NotNil(t, c.Take(nil, "foo", func(v interface{}) error {
+		return nil
+	}))
+	assert.NotNil(t, c.TakeWithExpire(nil, "foo", func(v interface{}, duration time.Duration) error {
+		return nil
+	}))
+}
+
 func calcEntropy(m map[int]int, total int) float64 {
 	var entropy float64
 

core/stores/cache/cachenode.go

@@ -1,13 +1,13 @@
 package cache
 
 import (
-	"encoding/json"
 	"errors"
 	"fmt"
 	"math/rand"
 	"sync"
 	"time"
 
+	"github.com/tal-tech/go-zero/core/jsonx"
 	"github.com/tal-tech/go-zero/core/logx"
 	"github.com/tal-tech/go-zero/core/mathx"
 	"github.com/tal-tech/go-zero/core/stat"
@@ -79,7 +79,7 @@ func (c cacheNode) SetCache(key string, v interface{}) error {
 }
 
 func (c cacheNode) SetCacheWithExpire(key string, v interface{}, expire time.Duration) error {
-	data, err := json.Marshal(v)
+	data, err := jsonx.Marshal(v)
 	if err != nil {
 		return err
 	}
@@ -168,7 +168,7 @@ func (c cacheNode) doTake(v interface{}, key string, query func(v interface{}) e
 			}
 		}
 
-		return json.Marshal(v)
+		return jsonx.Marshal(v)
 	})
 	if err != nil {
 		return err
@@ -181,11 +181,11 @@ func (c cacheNode) doTake(v interface{}, key string, query func(v interface{}) e
 		c.stat.IncrementHit()
 	}
 
-	return json.Unmarshal(val.([]byte), v)
+	return jsonx.Unmarshal(val.([]byte), v)
 }
 
 func (c cacheNode) processCache(key string, data string, v interface{}) error {
-	err := json.Unmarshal([]byte(data), v)
+	err := jsonx.Unmarshal([]byte(data), v)
 	if err == nil {
 		return nil
 	}

core/stores/cache/cachenode_test.go

@@ -2,7 +2,9 @@ package cache
 
 import (
 	"errors"
+	"fmt"
 	"math/rand"
+	"strconv"
 	"sync"
 	"testing"
 	"time"
@@ -13,8 +15,11 @@ import (
 	"github.com/tal-tech/go-zero/core/mathx"
 	"github.com/tal-tech/go-zero/core/stat"
 	"github.com/tal-tech/go-zero/core/stores/redis"
+	"github.com/tal-tech/go-zero/core/syncx"
 )
 
+var errTestNotFound = errors.New("not found")
+
 func init() {
 	logx.Disable()
 	stat.SetReporter(nil)
@@ -31,7 +36,7 @@ func TestCacheNode_DelCache(t *testing.T) {
 		lock:           new(sync.Mutex),
 		unstableExpiry: mathx.NewUnstable(expiryDeviation),
 		stat:           NewCacheStat("any"),
-		errNotFound:    errors.New("any"),
+		errNotFound:    errTestNotFound,
 	}
 	assert.Nil(t, cn.DelCache())
 	assert.Nil(t, cn.DelCache([]string{}...))
@@ -54,7 +59,7 @@ func TestCacheNode_InvalidCache(t *testing.T) {
 		lock:           new(sync.Mutex),
 		unstableExpiry: mathx.NewUnstable(expiryDeviation),
 		stat:           NewCacheStat("any"),
-		errNotFound:    errors.New("any"),
+		errNotFound:    errTestNotFound,
 	}
 	s.Set("any", "value")
 	var str string
@@ -63,3 +68,141 @@ func TestCacheNode_InvalidCache(t *testing.T) {
 	_, err = s.Get("any")
 	assert.Equal(t, miniredis.ErrKeyNotFound, err)
 }
+
+func TestCacheNode_Take(t *testing.T) {
+	s, err := miniredis.Run()
+	assert.Nil(t, err)
+	defer s.Close()
+
+	cn := cacheNode{
+		rds:            redis.NewRedis(s.Addr(), redis.NodeType),
+		r:              rand.New(rand.NewSource(time.Now().UnixNano())),
+		barrier:        syncx.NewSharedCalls(),
+		lock:           new(sync.Mutex),
+		unstableExpiry: mathx.NewUnstable(expiryDeviation),
+		stat:           NewCacheStat("any"),
+		errNotFound:    errTestNotFound,
+	}
+	var str string
+	err = cn.Take(&str, "any", func(v interface{}) error {
+		*v.(*string) = "value"
+		return nil
+	})
+	assert.Nil(t, err)
+	assert.Equal(t, "value", str)
+	assert.Nil(t, cn.GetCache("any", &str))
+	val, err := s.Get("any")
+	assert.Nil(t, err)
+	assert.Equal(t, `"value"`, val)
+}
+
+func TestCacheNode_TakeNotFound(t *testing.T) {
+	s, err := miniredis.Run()
+	assert.Nil(t, err)
+	defer s.Close()
+
+	cn := cacheNode{
+		rds:            redis.NewRedis(s.Addr(), redis.NodeType),
+		r:              rand.New(rand.NewSource(time.Now().UnixNano())),
+		barrier:        syncx.NewSharedCalls(),
+		lock:           new(sync.Mutex),
+		unstableExpiry: mathx.NewUnstable(expiryDeviation),
+		stat:           NewCacheStat("any"),
+		errNotFound:    errTestNotFound,
+	}
+	var str string
+	err = cn.Take(&str, "any", func(v interface{}) error {
+		return errTestNotFound
+	})
+	assert.Equal(t, errTestNotFound, err)
+	assert.Equal(t, errTestNotFound, cn.GetCache("any", &str))
+	val, err := s.Get("any")
+	assert.Nil(t, err)
+	assert.Equal(t, `*`, val)
+
+	s.Set("any", "*")
+	err = cn.Take(&str, "any", func(v interface{}) error {
+		return nil
+	})
+	assert.Equal(t, errTestNotFound, err)
+	assert.Equal(t, errTestNotFound, cn.GetCache("any", &str))
+
+	s.Del("any")
+	var errDummy = errors.New("dummy")
+	err = cn.Take(&str, "any", func(v interface{}) error {
+		return errDummy
+	})
+	assert.Equal(t, errDummy, err)
+}
+
+func TestCacheNode_TakeWithExpire(t *testing.T) {
+	s, err := miniredis.Run()
+	assert.Nil(t, err)
+	defer s.Close()
+
+	cn := cacheNode{
+		rds:            redis.NewRedis(s.Addr(), redis.NodeType),
+		r:              rand.New(rand.NewSource(time.Now().UnixNano())),
+		barrier:        syncx.NewSharedCalls(),
+		lock:           new(sync.Mutex),
+		unstableExpiry: mathx.NewUnstable(expiryDeviation),
+		stat:           NewCacheStat("any"),
+		errNotFound:    errors.New("any"),
+	}
+	var str string
+	err = cn.TakeWithExpire(&str, "any", func(v interface{}, expire time.Duration) error {
+		*v.(*string) = "value"
+		return nil
+	})
+	assert.Nil(t, err)
+	assert.Equal(t, "value", str)
+	assert.Nil(t, cn.GetCache("any", &str))
+	val, err := s.Get("any")
+	assert.Nil(t, err)
+	assert.Equal(t, `"value"`, val)
+}
+
+func TestCacheNode_String(t *testing.T) {
+	s, err := miniredis.Run()
+	assert.Nil(t, err)
+	defer s.Close()
+
+	cn := cacheNode{
+		rds:            redis.NewRedis(s.Addr(), redis.NodeType),
+		r:              rand.New(rand.NewSource(time.Now().UnixNano())),
+		barrier:        syncx.NewSharedCalls(),
+		lock:           new(sync.Mutex),
+		unstableExpiry: mathx.NewUnstable(expiryDeviation),
+		stat:           NewCacheStat("any"),
+		errNotFound:    errors.New("any"),
+	}
+	assert.Equal(t, s.Addr(), cn.String())
+}
+
+func TestCacheValueWithBigInt(t *testing.T) {
+	s, err := miniredis.Run()
+	if err != nil {
+		t.Error(err)
+	}
+	defer s.Close()
+
+	cn := cacheNode{
+		rds:            redis.NewRedis(s.Addr(), redis.NodeType),
+		r:              rand.New(rand.NewSource(time.Now().UnixNano())),
+		barrier:        syncx.NewSharedCalls(),
+		lock:           new(sync.Mutex),
+		unstableExpiry: mathx.NewUnstable(expiryDeviation),
+		stat:           NewCacheStat("any"),
+		errNotFound:    errors.New("any"),
+	}
+
+	const (
+		key         = "key"
+		value int64 = 323427211229009810
+	)
+
+	assert.Nil(t, cn.SetCache(key, value))
+	var val interface{}
+	assert.Nil(t, cn.GetCache(key, &val))
+	assert.Equal(t, strconv.FormatInt(value, 10), fmt.Sprintf("%v", val))
+}

core/stores/kv/store_test.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/alicebob/miniredis"
 	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/hash"
 	"github.com/tal-tech/go-zero/core/stores/cache"
 	"github.com/tal-tech/go-zero/core/stores/redis"
 	"github.com/tal-tech/go-zero/core/stringx"
@@ -15,6 +16,10 @@ var s1, _ = miniredis.Run()
 var s2, _ = miniredis.Run()
 
 func TestRedis_Exists(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Exists("foo")
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		ok, err := client.Exists("a")
 		assert.Nil(t, err)
@@ -27,6 +32,10 @@ func TestRedis_Exists(t *testing.T) {
 }
 
 func TestRedis_Eval(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Eval(`redis.call("EXISTS", KEYS[1])`, "key1")
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		_, err := client.Eval(`redis.call("EXISTS", KEYS[1])`, "notexist")
 		assert.Equal(t, redis.Nil, err)
@@ -41,6 +50,12 @@ func TestRedis_Eval(t *testing.T) {
 }
 
 func TestRedis_Hgetall(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	err := store.Hset("a", "aa", "aaa")
+	assert.NotNil(t, err)
+	_, err = store.Hgetall("a")
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		assert.Nil(t, client.Hset("a", "aa", "aaa"))
 		assert.Nil(t, client.Hset("a", "bb", "bbb"))
@@ -54,6 +69,10 @@ func TestRedis_Hgetall(t *testing.T) {
 }
 
 func TestRedis_Hvals(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Hvals("a")
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		assert.Nil(t, client.Hset("a", "aa", "aaa"))
 		assert.Nil(t, client.Hset("a", "bb", "bbb"))
@@ -64,6 +83,10 @@ func TestRedis_Hvals(t *testing.T) {
 }
 
 func TestRedis_Hsetnx(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Hsetnx("a", "dd", "ddd")
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		assert.Nil(t, client.Hset("a", "aa", "aaa"))
 		assert.Nil(t, client.Hset("a", "bb", "bbb"))
@@ -80,6 +103,12 @@ func TestRedis_Hsetnx(t *testing.T) {
 }
 
 func TestRedis_HdelHlen(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Hdel("a", "aa")
+	assert.NotNil(t, err)
+	_, err = store.Hlen("a")
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		assert.Nil(t, client.Hset("a", "aa", "aaa"))
 		assert.Nil(t, client.Hset("a", "bb", "bbb"))
@@ -96,6 +125,10 @@ func TestRedis_HdelHlen(t *testing.T) {
 }
 
 func TestRedis_HIncrBy(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Hincrby("key", "field", 3)
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		val, err := client.Hincrby("key", "field", 2)
 		assert.Nil(t, err)
@@ -107,6 +140,10 @@ func TestRedis_HIncrBy(t *testing.T) {
 }
 
 func TestRedis_Hkeys(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Hkeys("a")
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		assert.Nil(t, client.Hset("a", "aa", "aaa"))
 		assert.Nil(t, client.Hset("a", "bb", "bbb"))
@@ -117,6 +154,10 @@ func TestRedis_Hkeys(t *testing.T) {
 }
 
 func TestRedis_Hmget(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Hmget("a", "aa", "bb")
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		assert.Nil(t, client.Hset("a", "aa", "aaa"))
 		assert.Nil(t, client.Hset("a", "bb", "bbb"))
@@ -130,6 +171,12 @@ func TestRedis_Hmget(t *testing.T) {
 }
 
 func TestRedis_Hmset(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	err := store.Hmset("a", map[string]string{
+		"aa": "aaa",
+	})
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		assert.Nil(t, client.Hmset("a", map[string]string{
 			"aa": "aaa",
@@ -142,6 +189,10 @@ func TestRedis_Hmset(t *testing.T) {
 }
 
 func TestRedis_Incr(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Incr("a")
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		val, err := client.Incr("a")
 		assert.Nil(t, err)
@@ -153,6 +204,10 @@ func TestRedis_Incr(t *testing.T) {
 }
 
 func TestRedis_IncrBy(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Incrby("a", 2)
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		val, err := client.Incrby("a", 2)
 		assert.Nil(t, err)
@@ -164,6 +219,20 @@ func TestRedis_IncrBy(t *testing.T) {
 }
 
 func TestRedis_List(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Lpush("key", "value1", "value2")
+	assert.NotNil(t, err)
+	_, err = store.Rpush("key", "value3", "value4")
+	assert.NotNil(t, err)
+	_, err = store.Llen("key")
+	assert.NotNil(t, err)
+	_, err = store.Lrange("key", 0, 10)
+	assert.NotNil(t, err)
+	_, err = store.Lpop("key")
+	assert.NotNil(t, err)
+	_, err = store.Lrem("key", 0, "val")
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		val, err := client.Lpush("key", "value1", "value2")
 		assert.Nil(t, err)
@@ -202,6 +271,14 @@ func TestRedis_List(t *testing.T) {
 }
 
 func TestRedis_Persist(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Persist("key")
+	assert.NotNil(t, err)
+	err = store.Expire("key", 5)
+	assert.NotNil(t, err)
+	err = store.Expireat("key", time.Now().Unix()+5)
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		ok, err := client.Persist("key")
 		assert.Nil(t, err)
@@ -225,8 +302,16 @@ func TestRedis_Persist(t *testing.T) {
 }
 
 func TestRedis_Sscan(t *testing.T) {
+	key := "list"
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Sadd(key, nil)
+	assert.NotNil(t, err)
+	_, _, err = store.Sscan(key, 0, "", 100)
+	assert.NotNil(t, err)
+	_, err = store.Del(key)
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
-		key := "list"
 		var list []string
 		for i := 0; i < 1550; i++ {
 			list = append(list, stringx.Randn(i))
@@ -254,6 +339,20 @@ func TestRedis_Sscan(t *testing.T) {
 }
 
 func TestRedis_Set(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Scard("key")
+	assert.NotNil(t, err)
+	_, err = store.Sismember("key", 2)
+	assert.NotNil(t, err)
+	_, err = store.Srem("key", 3, 4)
+	assert.NotNil(t, err)
+	_, err = store.Smembers("key")
+	assert.NotNil(t, err)
+	_, err = store.Srandmember("key", 1)
+	assert.NotNil(t, err)
+	_, err = store.Spop("key")
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		num, err := client.Sadd("key", 1, 2, 3, 4)
 		assert.Nil(t, err)
@@ -290,6 +389,14 @@ func TestRedis_Set(t *testing.T) {
 }
 
 func TestRedis_SetGetDel(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	err := store.Set("hello", "world")
+	assert.NotNil(t, err)
+	_, err = store.Get("hello")
+	assert.NotNil(t, err)
+	_, err = store.Del("hello")
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		err := client.Set("hello", "world")
 		assert.Nil(t, err)
@@ -303,6 +410,16 @@ func TestRedis_SetGetDel(t *testing.T) {
 }
 
 func TestRedis_SetExNx(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	err := store.Setex("hello", "world", 5)
+	assert.NotNil(t, err)
+	_, err = store.Setnx("newhello", "newworld")
+	assert.NotNil(t, err)
+	_, err = store.Ttl("hello")
+	assert.NotNil(t, err)
+	_, err = store.SetnxEx("newhello", "newworld", 5)
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		err := client.Setex("hello", "world", 5)
 		assert.Nil(t, err)
@@ -337,6 +454,16 @@ func TestRedis_SetExNx(t *testing.T) {
 }
 
 func TestRedis_SetGetDelHashField(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	err := store.Hset("key", "field", "value")
+	assert.NotNil(t, err)
+	_, err = store.Hget("key", "field")
+	assert.NotNil(t, err)
+	_, err = store.Hexists("key", "field")
+	assert.NotNil(t, err)
+	_, err = store.Hdel("key", "field")
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		err := client.Hset("key", "field", "value")
 		assert.Nil(t, err)
@@ -356,6 +483,48 @@ func TestRedis_SetGetDelHashField(t *testing.T) {
 }
 
 func TestRedis_SortedSet(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Zadd("key", 1, "value1")
+	assert.NotNil(t, err)
+	_, err = store.Zscore("key", "value1")
+	assert.NotNil(t, err)
+	_, err = store.Zcount("key", 6, 7)
+	assert.NotNil(t, err)
+	_, err = store.Zincrby("key", 3, "value1")
+	assert.NotNil(t, err)
+	_, err = store.Zrank("key", "value2")
+	assert.NotNil(t, err)
+	_, err = store.Zrem("key", "value2", "value3")
+	assert.NotNil(t, err)
+	_, err = store.Zremrangebyscore("key", 6, 7)
+	assert.NotNil(t, err)
+	_, err = store.Zremrangebyrank("key", 1, 2)
+	assert.NotNil(t, err)
+	_, err = store.Zcard("key")
+	assert.NotNil(t, err)
+	_, err = store.Zrange("key", 0, -1)
+	assert.NotNil(t, err)
+	_, err = store.Zrevrange("key", 0, -1)
+	assert.NotNil(t, err)
+	_, err = store.ZrangeWithScores("key", 0, -1)
+	assert.NotNil(t, err)
+	_, err = store.ZrangebyscoreWithScores("key", 5, 8)
+	assert.NotNil(t, err)
+	_, err = store.ZrangebyscoreWithScoresAndLimit("key", 5, 8, 1, 1)
+	assert.NotNil(t, err)
+	_, err = store.ZrevrangebyscoreWithScores("key", 5, 8)
+	assert.NotNil(t, err)
+	_, err = store.ZrevrangebyscoreWithScoresAndLimit("key", 5, 8, 1, 1)
+	assert.NotNil(t, err)
+	_, err = store.Zadds("key", redis.Pair{
+		Key:   "value2",
+		Score: 6,
+	}, redis.Pair{
+		Key:   "value3",
+		Score: 7,
+	})
+	assert.NotNil(t, err)
+
 	runOnCluster(t, func(client Store) {
 		ok, err := client.Zadd("key", 1, "value1")
 		assert.Nil(t, err)
@@ -471,6 +640,30 @@ func TestRedis_SortedSet(t *testing.T) {
 				Score: 5,
 			},
 		}, pairs)
+		val, err = client.Zadds("key", redis.Pair{
+			Key:   "value2",
+			Score: 6,
+		}, redis.Pair{
+			Key:   "value3",
+			Score: 7,
+		})
+		assert.Nil(t, err)
+		assert.Equal(t, int64(2), val)
+	})
+}
+
+func TestRedis_HyperLogLog(t *testing.T) {
+	store := clusterStore{dispatcher: hash.NewConsistentHash()}
+	_, err := store.Pfadd("key")
+	assert.NotNil(t, err)
+	_, err = store.Pfcount("key")
+	assert.NotNil(t, err)
+
+	runOnCluster(t, func(cluster Store) {
+		_, err := cluster.Pfadd("key")
+		assert.NotNil(t, err)
+		_, err = cluster.Pfcount("key")
+		assert.NotNil(t, err)
 	})
 }
 

core/stores/mongo/collection.go

@@ -7,6 +7,7 @@ import (
 	"github.com/globalsign/mgo"
 	"github.com/tal-tech/go-zero/core/breaker"
 	"github.com/tal-tech/go-zero/core/logx"
+	"github.com/tal-tech/go-zero/core/stores/mongo/internal"
 	"github.com/tal-tech/go-zero/core/timex"
 )
 
@@ -29,8 +30,9 @@ type (
 	}
 
 	decoratedCollection struct {
-		*mgo.Collection
-		brk breaker.Breaker
+		name       string
+		collection internal.MgoCollection
+		brk        breaker.Breaker
 	}
 
 	keepablePromise struct {
@@ -41,7 +43,8 @@ type (
 
 func newCollection(collection *mgo.Collection) Collection {
 	return &decoratedCollection{
-		Collection: collection,
+		name:       collection.FullName,
+		collection: collection,
 		brk:        breaker.NewBreaker(),
 	}
 }
@@ -54,7 +57,7 @@ func (c *decoratedCollection) Find(query interface{}) Query {
 
 	startTime := timex.Now()
 	return promisedQuery{
-		Query: c.Collection.Find(query),
+		Query: c.collection.Find(query),
 		promise: keepablePromise{
 			promise: promise,
 			log: func(err error) {
@@ -73,7 +76,7 @@ func (c *decoratedCollection) FindId(id interface{}) Query {
 
 	startTime := timex.Now()
 	return promisedQuery{
-		Query: c.Collection.FindId(id),
+		Query: c.collection.FindId(id),
 		promise: keepablePromise{
 			promise: promise,
 			log: func(err error) {
@@ -92,7 +95,7 @@ func (c *decoratedCollection) Insert(docs ...interface{}) (err error) {
 			c.logDuration("insert", duration, err, docs...)
 		}()
 
-		return c.Collection.Insert(docs...)
+		return c.collection.Insert(docs...)
 	}, acceptable)
 }
 
@@ -104,7 +107,7 @@ func (c *decoratedCollection) Pipe(pipeline interface{}) Pipe {
 
 	startTime := timex.Now()
 	return promisedPipe{
-		Pipe: c.Collection.Pipe(pipeline),
+		Pipe: c.collection.Pipe(pipeline),
 		promise: keepablePromise{
 			promise: promise,
 			log: func(err error) {
@@ -123,7 +126,7 @@ func (c *decoratedCollection) Remove(selector interface{}) (err error) {
 			c.logDuration("remove", duration, err, selector)
 		}()
 
-		return c.Collection.Remove(selector)
+		return c.collection.Remove(selector)
 	}, acceptable)
 }
 
@@ -135,7 +138,7 @@ func (c *decoratedCollection) RemoveAll(selector interface{}) (info *mgo.ChangeI
 			c.logDuration("removeAll", duration, err, selector)
 		}()
 
-		info, err = c.Collection.RemoveAll(selector)
+		info, err = c.collection.RemoveAll(selector)
 		return err
 	}, acceptable)
 
@@ -150,7 +153,7 @@ func (c *decoratedCollection) RemoveId(id interface{}) (err error) {
 			c.logDuration("removeId", duration, err, id)
 		}()
 
-		return c.Collection.RemoveId(id)
+		return c.collection.RemoveId(id)
 	}, acceptable)
 }
 
@@ -162,7 +165,7 @@ func (c *decoratedCollection) Update(selector, update interface{}) (err error) {
 			c.logDuration("update", duration, err, selector, update)
 		}()
 
-		return c.Collection.Update(selector, update)
+		return c.collection.Update(selector, update)
 	}, acceptable)
 }
 
@@ -174,7 +177,7 @@ func (c *decoratedCollection) UpdateId(id, update interface{}) (err error) {
 			c.logDuration("updateId", duration, err, id, update)
 		}()
 
-		return c.Collection.UpdateId(id, update)
+		return c.collection.UpdateId(id, update)
 	}, acceptable)
 }
 
@@ -186,7 +189,7 @@ func (c *decoratedCollection) Upsert(selector, update interface{}) (info *mgo.Ch
 			c.logDuration("upsert", duration, err, selector, update)
 		}()
 
-		info, err = c.Collection.Upsert(selector, update)
+		info, err = c.collection.Upsert(selector, update)
 		return err
 	}, acceptable)
 
@@ -200,17 +203,17 @@ func (c *decoratedCollection) logDuration(method string, duration time.Duration,
 	} else if err != nil {
 		if duration > slowThreshold {
 			logx.WithDuration(duration).Slowf("[MONGO] mongo(%s) - slowcall - %s - fail(%s) - %s",
-				c.FullName, method, err.Error(), string(content))
+				c.name, method, err.Error(), string(content))
 		} else {
 			logx.WithDuration(duration).Infof("mongo(%s) - %s - fail(%s) - %s",
-				c.FullName, method, err.Error(), string(content))
+				c.name, method, err.Error(), string(content))
 		}
 	} else {
 		if duration > slowThreshold {
 			logx.WithDuration(duration).Slowf("[MONGO] mongo(%s) - slowcall - %s - ok - %s",
-				c.FullName, method, string(content))
+				c.name, method, string(content))
 		} else {
-			logx.WithDuration(duration).Infof("mongo(%s) - %s - ok - %s", c.FullName, method, string(content))
+			logx.WithDuration(duration).Infof("mongo(%s) - %s - ok - %s", c.name, method, string(content))
 		}
 	}
 }

core/stores/mongo/collection_test.go

@@ -5,10 +5,20 @@ import (
 	"testing"
 
 	"github.com/globalsign/mgo"
+	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/breaker"
+	"github.com/tal-tech/go-zero/core/logx"
+	"github.com/tal-tech/go-zero/core/stores/mongo/internal"
 	"github.com/tal-tech/go-zero/core/stringx"
 )
 
+var errDummy = errors.New("dummy")
+
+func init() {
+	logx.Disable()
+}
+
 func TestKeepPromise_accept(t *testing.T) {
 	p := new(mockPromise)
 	kp := keepablePromise{
@@ -56,6 +66,206 @@ func TestKeepPromise_keep(t *testing.T) {
 	}
 }
 
+func TestNewCollection(t *testing.T) {
+	col := newCollection(&mgo.Collection{
+		Database: nil,
+		Name:     "foo",
+		FullName: "bar",
+	})
+	assert.Equal(t, "bar", col.(*decoratedCollection).name)
+}
+
+func TestCollectionFind(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	var query mgo.Query
+	col := internal.NewMockMgoCollection(ctrl)
+	col.EXPECT().Find(gomock.Any()).Return(&query)
+	c := decoratedCollection{
+		collection: col,
+		brk:        breaker.NewBreaker(),
+	}
+	actual := c.Find(nil)
+	switch v := actual.(type) {
+	case promisedQuery:
+		assert.Equal(t, &query, v.Query)
+		assert.Equal(t, errDummy, v.promise.keep(errDummy))
+	default:
+		t.Fail()
+	}
+	c.brk = new(dropBreaker)
+	actual = c.Find(nil)
+	assert.Equal(t, rejectedQuery{}, actual)
+}
+
+func TestCollectionFindId(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	var query mgo.Query
+	col := internal.NewMockMgoCollection(ctrl)
+	col.EXPECT().FindId(gomock.Any()).Return(&query)
+	c := decoratedCollection{
+		collection: col,
+		brk:        breaker.NewBreaker(),
+	}
+	actual := c.FindId(nil)
+	switch v := actual.(type) {
+	case promisedQuery:
+		assert.Equal(t, &query, v.Query)
+		assert.Equal(t, errDummy, v.promise.keep(errDummy))
+	default:
+		t.Fail()
+	}
+	c.brk = new(dropBreaker)
+	actual = c.FindId(nil)
+	assert.Equal(t, rejectedQuery{}, actual)
+}
+
+func TestCollectionInsert(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	col := internal.NewMockMgoCollection(ctrl)
+	col.EXPECT().Insert(nil, nil).Return(errDummy)
+	c := decoratedCollection{
+		collection: col,
+		brk:        breaker.NewBreaker(),
+	}
+	err := c.Insert(nil, nil)
+	assert.Equal(t, errDummy, err)
+	c.brk = new(dropBreaker)
+	err = c.Insert(nil, nil)
+	assert.Equal(t, errDummy, err)
+}
+
+func TestCollectionPipe(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	var pipe mgo.Pipe
+	col := internal.NewMockMgoCollection(ctrl)
+	col.EXPECT().Pipe(gomock.Any()).Return(&pipe)
+	c := decoratedCollection{
+		collection: col,
+		brk:        breaker.NewBreaker(),
+	}
+	actual := c.Pipe(nil)
+	switch v := actual.(type) {
+	case promisedPipe:
+		assert.Equal(t, &pipe, v.Pipe)
+		assert.Equal(t, errDummy, v.promise.keep(errDummy))
+	default:
+		t.Fail()
+	}
+	c.brk = new(dropBreaker)
+	actual = c.Pipe(nil)
+	assert.Equal(t, rejectedPipe{}, actual)
+}
+
+func TestCollectionRemove(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	col := internal.NewMockMgoCollection(ctrl)
+	col.EXPECT().Remove(gomock.Any()).Return(errDummy)
+	c := decoratedCollection{
+		collection: col,
+		brk:        breaker.NewBreaker(),
+	}
+	err := c.Remove(nil)
+	assert.Equal(t, errDummy, err)
+	c.brk = new(dropBreaker)
+	err = c.Remove(nil)
+	assert.Equal(t, errDummy, err)
+}
+
+func TestCollectionRemoveAll(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	col := internal.NewMockMgoCollection(ctrl)
+	col.EXPECT().RemoveAll(gomock.Any()).Return(nil, errDummy)
+	c := decoratedCollection{
+		collection: col,
+		brk:        breaker.NewBreaker(),
+	}
+	_, err := c.RemoveAll(nil)
+	assert.Equal(t, errDummy, err)
+	c.brk = new(dropBreaker)
+	_, err = c.RemoveAll(nil)
+	assert.Equal(t, errDummy, err)
+}
+
+func TestCollectionRemoveId(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	col := internal.NewMockMgoCollection(ctrl)
+	col.EXPECT().RemoveId(gomock.Any()).Return(errDummy)
+	c := decoratedCollection{
+		collection: col,
+		brk:        breaker.NewBreaker(),
+	}
+	err := c.RemoveId(nil)
+	assert.Equal(t, errDummy, err)
+	c.brk = new(dropBreaker)
+	err = c.RemoveId(nil)
+	assert.Equal(t, errDummy, err)
+}
+
+func TestCollectionUpdate(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	col := internal.NewMockMgoCollection(ctrl)
+	col.EXPECT().Update(gomock.Any(), gomock.Any()).Return(errDummy)
+	c := decoratedCollection{
+		collection: col,
+		brk:        breaker.NewBreaker(),
+	}
+	err := c.Update(nil, nil)
+	assert.Equal(t, errDummy, err)
+	c.brk = new(dropBreaker)
+	err = c.Update(nil, nil)
+	assert.Equal(t, errDummy, err)
+}
+
+func TestCollectionUpdateId(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	col := internal.NewMockMgoCollection(ctrl)
+	col.EXPECT().UpdateId(gomock.Any(), gomock.Any()).Return(errDummy)
+	c := decoratedCollection{
+		collection: col,
+		brk:        breaker.NewBreaker(),
+	}
+	err := c.UpdateId(nil, nil)
+	assert.Equal(t, errDummy, err)
+	c.brk = new(dropBreaker)
+	err = c.UpdateId(nil, nil)
+	assert.Equal(t, errDummy, err)
+}
+
+func TestCollectionUpsert(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	col := internal.NewMockMgoCollection(ctrl)
+	col.EXPECT().Upsert(gomock.Any(), gomock.Any()).Return(nil, errDummy)
+	c := decoratedCollection{
+		collection: col,
+		brk:        breaker.NewBreaker(),
+	}
+	_, err := c.Upsert(nil, nil)
+	assert.Equal(t, errDummy, err)
+	c.brk = new(dropBreaker)
+	_, err = c.Upsert(nil, nil)
+	assert.Equal(t, errDummy, err)
+}
+
 type mockPromise struct {
 	accepted bool
 	reason   string
@@ -68,3 +278,31 @@ func (p *mockPromise) Accept() {
 func (p *mockPromise) Reject(reason string) {
 	p.reason = reason
 }
+
+type dropBreaker struct {
+}
+
+func (d *dropBreaker) Name() string {
+	return "dummy"
+}
+
+func (d *dropBreaker) Allow() (breaker.Promise, error) {
+	return nil, errDummy
+}
+
+func (d *dropBreaker) Do(req func() error) error {
+	return nil
+}
+
+func (d *dropBreaker) DoWithAcceptable(req func() error, acceptable breaker.Acceptable) error {
+	return errDummy
+}
+
+func (d *dropBreaker) DoWithFallback(req func() error, fallback func(err error) error) error {
+	return nil
+}
+
+func (d *dropBreaker) DoWithFallbackAcceptable(req func() error, fallback func(err error) error,
+	acceptable breaker.Acceptable) error {
+	return nil
+}

core/stores/mongo/internal/collection.go

@@ -0,0 +1,17 @@
+//go:generate mockgen -package internal -destination collection_mock.go -source collection.go
+package internal
+
+import "github.com/globalsign/mgo"
+
+type MgoCollection interface {
+	Find(query interface{}) *mgo.Query
+	FindId(id interface{}) *mgo.Query
+	Insert(docs ...interface{}) error
+	Pipe(pipeline interface{}) *mgo.Pipe
+	Remove(selector interface{}) error
+	RemoveAll(selector interface{}) (*mgo.ChangeInfo, error)
+	RemoveId(id interface{}) error
+	Update(selector, update interface{}) error
+	UpdateId(id, update interface{}) error
+	Upsert(selector, update interface{}) (*mgo.ChangeInfo, error)
+}

core/stores/mongo/internal/collection_mock.go

@@ -0,0 +1,180 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: collection.go
+
+// Package internal is a generated GoMock package.
+package internal
+
+import (
+	mgo "github.com/globalsign/mgo"
+	gomock "github.com/golang/mock/gomock"
+	reflect "reflect"
+)
+
+// MockMgoCollection is a mock of MgoCollection interface
+type MockMgoCollection struct {
+	ctrl     *gomock.Controller
+	recorder *MockMgoCollectionMockRecorder
+}
+
+// MockMgoCollectionMockRecorder is the mock recorder for MockMgoCollection
+type MockMgoCollectionMockRecorder struct {
+	mock *MockMgoCollection
+}
+
+// NewMockMgoCollection creates a new mock instance
+func NewMockMgoCollection(ctrl *gomock.Controller) *MockMgoCollection {
+	mock := &MockMgoCollection{ctrl: ctrl}
+	mock.recorder = &MockMgoCollectionMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use
+func (m *MockMgoCollection) EXPECT() *MockMgoCollectionMockRecorder {
+	return m.recorder
+}
+
+// Find mocks base method
+func (m *MockMgoCollection) Find(query interface{}) *mgo.Query {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Find", query)
+	ret0, _ := ret[0].(*mgo.Query)
+	return ret0
+}
+
+// Find indicates an expected call of Find
+func (mr *MockMgoCollectionMockRecorder) Find(query interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Find", reflect.TypeOf((*MockMgoCollection)(nil).Find), query)
+}
+
+// FindId mocks base method
+func (m *MockMgoCollection) FindId(id interface{}) *mgo.Query {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "FindId", id)
+	ret0, _ := ret[0].(*mgo.Query)
+	return ret0
+}
+
+// FindId indicates an expected call of FindId
+func (mr *MockMgoCollectionMockRecorder) FindId(id interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "FindId", reflect.TypeOf((*MockMgoCollection)(nil).FindId), id)
+}
+
+// Insert mocks base method
+func (m *MockMgoCollection) Insert(docs ...interface{}) error {
+	m.ctrl.T.Helper()
+	varargs := []interface{}{}
+	for _, a := range docs {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "Insert", varargs...)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Insert indicates an expected call of Insert
+func (mr *MockMgoCollectionMockRecorder) Insert(docs ...interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Insert", reflect.TypeOf((*MockMgoCollection)(nil).Insert), docs...)
+}
+
+// Pipe mocks base method
+func (m *MockMgoCollection) Pipe(pipeline interface{}) *mgo.Pipe {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Pipe", pipeline)
+	ret0, _ := ret[0].(*mgo.Pipe)
+	return ret0
+}
+
+// Pipe indicates an expected call of Pipe
+func (mr *MockMgoCollectionMockRecorder) Pipe(pipeline interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Pipe", reflect.TypeOf((*MockMgoCollection)(nil).Pipe), pipeline)
+}
+
+// Remove mocks base method
+func (m *MockMgoCollection) Remove(selector interface{}) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Remove", selector)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Remove indicates an expected call of Remove
+func (mr *MockMgoCollectionMockRecorder) Remove(selector interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Remove", reflect.TypeOf((*MockMgoCollection)(nil).Remove), selector)
+}
+
+// RemoveAll mocks base method
+func (m *MockMgoCollection) RemoveAll(selector interface{}) (*mgo.ChangeInfo, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "RemoveAll", selector)
+	ret0, _ := ret[0].(*mgo.ChangeInfo)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// RemoveAll indicates an expected call of RemoveAll
+func (mr *MockMgoCollectionMockRecorder) RemoveAll(selector interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RemoveAll", reflect.TypeOf((*MockMgoCollection)(nil).RemoveAll), selector)
+}
+
+// RemoveId mocks base method
+func (m *MockMgoCollection) RemoveId(id interface{}) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "RemoveId", id)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// RemoveId indicates an expected call of RemoveId
+func (mr *MockMgoCollectionMockRecorder) RemoveId(id interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RemoveId", reflect.TypeOf((*MockMgoCollection)(nil).RemoveId), id)
+}
+
+// Update mocks base method
+func (m *MockMgoCollection) Update(selector, update interface{}) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Update", selector, update)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Update indicates an expected call of Update
+func (mr *MockMgoCollectionMockRecorder) Update(selector, update interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Update", reflect.TypeOf((*MockMgoCollection)(nil).Update), selector, update)
+}
+
+// UpdateId mocks base method
+func (m *MockMgoCollection) UpdateId(id, update interface{}) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "UpdateId", id, update)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// UpdateId indicates an expected call of UpdateId
+func (mr *MockMgoCollectionMockRecorder) UpdateId(id, update interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateId", reflect.TypeOf((*MockMgoCollection)(nil).UpdateId), id, update)
+}
+
+// Upsert mocks base method
+func (m *MockMgoCollection) Upsert(selector, update interface{}) (*mgo.ChangeInfo, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Upsert", selector, update)
+	ret0, _ := ret[0].(*mgo.ChangeInfo)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// Upsert indicates an expected call of Upsert
+func (mr *MockMgoCollectionMockRecorder) Upsert(selector, update interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Upsert", reflect.TypeOf((*MockMgoCollection)(nil).Upsert), selector, update)
+}

core/stores/sqlc/cachedsql.go

@@ -83,10 +83,6 @@ func (cc CachedConn) QueryRowIndex(v interface{}, key string, keyer func(primary
 	var primaryKey interface{}
 	var found bool
 
-	// if don't use convert numeric primary key into int64,
-	// then it will be represented as scientific notion, like 2e6
-	// which will make the cache doesn't match with the previous insert one
-	keyer = floatKeyer(keyer)
 	if err := cc.cache.TakeWithExpire(&primaryKey, key, func(val interface{}, expire time.Duration) (err error) {
 		primaryKey, err = indexQuery(cc.db, v)
 		if err != nil {
@@ -124,16 +120,3 @@ func (cc CachedConn) SetCache(key string, v interface{}) error {
 func (cc CachedConn) Transact(fn func(sqlx.Session) error) error {
 	return cc.db.Transact(fn)
 }
-
-func floatKeyer(fn func(interface{}) string) func(interface{}) string {
-	return func(primary interface{}) string {
-		switch v := primary.(type) {
-		case float32:
-			return fn(int64(v))
-		case float64:
-			return fn(int64(v))
-		default:
-			return fn(primary)
-		}
-	}
-}

core/stores/sqlc/cachedsql_test.go

@@ -79,9 +79,29 @@ func TestCachedConn_QueryRowIndex_NoCache(t *testing.T) {
 	}
 
 	r := redis.NewRedis(s.Addr(), redis.NodeType)
-	c := NewNodeConn(dummySqlConn{}, r, cache.WithExpiry(time.Second*10))
+	c := NewConn(dummySqlConn{}, cache.CacheConf{
+		{
+			RedisConf: redis.RedisConf{
+				Host: s.Addr(),
+				Type: redis.NodeType,
+			},
+			Weight: 100,
+		},
+	}, cache.WithExpiry(time.Second*10))
 
 	var str string
+	err = c.QueryRowIndex(&str, "index", func(s interface{}) string {
+		return fmt.Sprintf("%s/1234", s)
+	}, func(conn sqlx.SqlConn, v interface{}) (interface{}, error) {
+		*v.(*string) = "zero"
+		return "primary", errors.New("foo")
+	}, func(conn sqlx.SqlConn, v, pri interface{}) error {
+		assert.Equal(t, "primary", pri)
+		*v.(*string) = "xin"
+		return nil
+	})
+	assert.NotNil(t, err)
+
 	err = c.QueryRowIndex(&str, "index", func(s interface{}) string {
 		return fmt.Sprintf("%s/1234", s)
 	}, func(conn sqlx.SqlConn, v interface{}) (interface{}, error) {
@@ -500,6 +520,10 @@ func TestCachedConnExecDropCache(t *testing.T) {
 	assert.True(t, conn.execValue)
 	_, err = s.Get(key)
 	assert.Exactly(t, miniredis.ErrKeyNotFound, err)
+	_, err = c.Exec(func(conn sqlx.SqlConn) (result sql.Result, e error) {
+		return nil, errors.New("foo")
+	}, key)
+	assert.NotNil(t, err)
 }
 
 func TestCachedConnExecDropCacheFailed(t *testing.T) {
@@ -570,20 +594,6 @@ func TestQueryRowNoCache(t *testing.T) {
 	assert.True(t, ran)
 }
 
-func TestFloatKeyer(t *testing.T) {
-	primaries := []interface{}{
-		float32(1),
-		float64(1),
-	}
-
-	for _, primary := range primaries {
-		val := floatKeyer(func(i interface{}) string {
-			return fmt.Sprint(i)
-		})(primary)
-		assert.Equal(t, "1", val)
-	}
-}
-
 func resetStats() {
 	atomic.StoreUint64(&stats.Total, 0)
 	atomic.StoreUint64(&stats.Hit, 0)

core/stores/sqlx/bulkinserter_test.go

@@ -2,6 +2,7 @@ package sqlx
 
 import (
 	"database/sql"
+	"errors"
 	"strconv"
 	"testing"
 
@@ -11,14 +12,15 @@ import (
 )
 
 type mockedConn struct {
-	query string
-	args  []interface{}
+	query   string
+	args    []interface{}
+	execErr error
 }
 
 func (c *mockedConn) Exec(query string, args ...interface{}) (sql.Result, error) {
 	c.query = query
 	c.args = args
-	return nil, nil
+	return nil, c.execErr
 }
 
 func (c *mockedConn) Prepare(query string) (StmtSession, error) {
@@ -68,9 +70,12 @@ func TestBulkInserterSuffix(t *testing.T) {
 		inserter, err := NewBulkInserter(&conn, `INSERT INTO classroom_dau(classroom, user, count) VALUES`+
 			`(?, ?, ?) ON DUPLICATE KEY UPDATE is_overtime=VALUES(is_overtime)`)
 		assert.Nil(t, err)
+		assert.Nil(t, inserter.UpdateStmt(`INSERT INTO classroom_dau(classroom, user, count) VALUES`+
+			`(?, ?, ?) ON DUPLICATE KEY UPDATE is_overtime=VALUES(is_overtime)`))
 		for i := 0; i < 5; i++ {
 			assert.Nil(t, inserter.Insert("class_"+strconv.Itoa(i), "user_"+strconv.Itoa(i), i))
 		}
+		inserter.SetResultHandler(func(result sql.Result, err error) {})
 		inserter.Flush()
 		assert.Equal(t, `INSERT INTO classroom_dau(classroom, user, count) VALUES `+
 			`('class_0', 'user_0', 0), ('class_1', 'user_1', 1), ('class_2', 'user_2', 2), `+
@@ -80,6 +85,33 @@ func TestBulkInserterSuffix(t *testing.T) {
 	})
 }
 
+func TestBulkInserterBadStatement(t *testing.T) {
+	runSqlTest(t, func(db *sql.DB, mock sqlmock.Sqlmock) {
+		var conn mockedConn
+		_, err := NewBulkInserter(&conn, "foo")
+		assert.NotNil(t, err)
+	})
+}
+
+func TestBulkInserter_Update(t *testing.T) {
+	conn := mockedConn{
+		execErr: errors.New("foo"),
+	}
+	_, err := NewBulkInserter(&conn, `INSERT INTO classroom_dau(classroom, user, count) VALUES()`)
+	assert.NotNil(t, err)
+	_, err = NewBulkInserter(&conn, `INSERT INTO classroom_dau(classroom, user, count) VALUES(?)`)
+	assert.NotNil(t, err)
+	inserter, err := NewBulkInserter(&conn, `INSERT INTO classroom_dau(classroom, user, count) VALUES(?, ?, ?)`)
+	assert.Nil(t, err)
+	inserter.inserter.Execute([]string{"bar"})
+	inserter.SetResultHandler(func(result sql.Result, err error) {
+	})
+	inserter.UpdateOrDelete(func() {})
+	inserter.inserter.Execute([]string(nil))
+	assert.NotNil(t, inserter.UpdateStmt("foo"))
+	assert.NotNil(t, inserter.Insert("foo", "bar"))
+}
+
 func runSqlTest(t *testing.T, fn func(db *sql.DB, mock sqlmock.Sqlmock)) {
 	logx.Disable()
 

core/stores/sqlx/orm_test.go

@@ -2,6 +2,7 @@ package sqlx
 
 import (
 	"database/sql"
+	"errors"
 	"testing"
 
 	"github.com/DATA-DOG/go-sqlmock"
@@ -256,24 +257,6 @@ func TestUnmarshalRowStructWithTagsWrongColumns(t *testing.T) {
 	})
 }
 
-func TestUnmarshalRowStructWithTagsPtr(t *testing.T) {
-	var value = new(struct {
-		Age  *int   `db:"age"`
-		Name string `db:"name"`
-	})
-
-	runOrmTest(t, func(db *sql.DB, mock sqlmock.Sqlmock) {
-		rs := sqlmock.NewRows([]string{"name", "age"}).FromCSVString("liao,5")
-		mock.ExpectQuery("select (.+) from users where user=?").WithArgs("anyone").WillReturnRows(rs)
-
-		assert.Nil(t, query(db, func(rows *sql.Rows) error {
-			return unmarshalRow(value, rows, true)
-		}, "select name, age from users where user=?", "anyone"))
-		assert.Equal(t, "liao", value.Name)
-		assert.Equal(t, 5, *value.Age)
-	})
-}
-
 func TestUnmarshalRowsBool(t *testing.T) {
 	runOrmTest(t, func(db *sql.DB, mock sqlmock.Sqlmock) {
 		var expect = []bool{true, false}
@@ -1001,6 +984,62 @@ func TestCommonSqlConn_QueryRowOptional(t *testing.T) {
 	})
 }
 
+func TestUnmarshalRowError(t *testing.T) {
+	tests := []struct {
+		name     string
+		colErr   error
+		scanErr  error
+		err      error
+		next     int
+		validate func(err error)
+	}{
+		{
+			name: "with error",
+			err:  errors.New("foo"),
+			validate: func(err error) {
+				assert.NotNil(t, err)
+			},
+		},
+		{
+			name: "without next",
+			validate: func(err error) {
+				assert.Equal(t, ErrNotFound, err)
+			},
+		},
+		{
+			name:    "with error",
+			scanErr: errors.New("foo"),
+			next:    1,
+			validate: func(err error) {
+				assert.Equal(t, ErrNotFound, err)
+			},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			runOrmTest(t, func(db *sql.DB, mock sqlmock.Sqlmock) {
+				rs := sqlmock.NewRows([]string{"age"}).FromCSVString("5")
+				mock.ExpectQuery("select (.+) from users where user=?").WithArgs(
+					"anyone").WillReturnRows(rs)
+
+				var r struct {
+					User string `db:"user"`
+					Age  int    `db:"age"`
+				}
+				test.validate(query(db, func(rows *sql.Rows) error {
+					scanner := mockedScanner{
+						colErr:  test.colErr,
+						scanErr: test.scanErr,
+						err:     test.err,
+					}
+					return unmarshalRow(&r, &scanner, false)
+				}, "select age from users where user=?", "anyone"))
+			})
+		})
+	}
+}
+
 func runOrmTest(t *testing.T, fn func(db *sql.DB, mock sqlmock.Sqlmock)) {
 	logx.Disable()
 
@@ -1016,3 +1055,30 @@ func runOrmTest(t *testing.T, fn func(db *sql.DB, mock sqlmock.Sqlmock)) {
 		t.Errorf("there were unfulfilled expectations: %s", err)
 	}
 }
+
+type mockedScanner struct {
+	colErr  error
+	scanErr error
+	err     error
+	next    int
+}
+
+func (m *mockedScanner) Columns() ([]string, error) {
+	return nil, m.colErr
+}
+
+func (m *mockedScanner) Err() error {
+	return m.err
+}
+
+func (m *mockedScanner) Next() bool {
+	if m.next > 0 {
+		m.next--
+		return true
+	}
+	return false
+}
+
+func (m *mockedScanner) Scan(v ...interface{}) error {
+	return m.scanErr
+}

core/syncx/limit.go

@@ -6,18 +6,22 @@ import (
 	"github.com/tal-tech/go-zero/core/lang"
 )
 
-var ErrReturn = errors.New("discarding limited token, resource pool is full, someone returned multiple times")
+// ErrLimitReturn indicates that the more than borrowed elements were returned.
+var ErrLimitReturn = errors.New("discarding limited token, resource pool is full, someone returned multiple times")
 
+// Limit controls the concurrent requests.
 type Limit struct {
 	pool chan lang.PlaceholderType
 }
 
+// NewLimit creates a Limit that can borrow n elements from it concurrently.
 func NewLimit(n int) Limit {
 	return Limit{
 		pool: make(chan lang.PlaceholderType, n),
 	}
 }
 
+// Borrow borrows an element from Limit in blocking mode.
 func (l Limit) Borrow() {
 	l.pool <- lang.Placeholder
 }
@@ -28,10 +32,12 @@ func (l Limit) Return() error {
 	case <-l.pool:
 		return nil
 	default:
-		return ErrReturn
+		return ErrLimitReturn
 	}
 }
 
+// TryBorrow tries to borrow an element from Limit, in non-blocking mode.
+// If success, true returned, false for otherwise.
 func (l Limit) TryBorrow() bool {
 	select {
 	case l.pool <- lang.Placeholder:

core/syncx/limit_test.go

@@ -13,5 +13,5 @@ func TestLimit(t *testing.T) {
 	assert.False(t, limit.TryBorrow())
 	assert.Nil(t, limit.Return())
 	assert.Nil(t, limit.Return())
-	assert.Equal(t, ErrReturn, limit.Return())
+	assert.Equal(t, ErrLimitReturn, limit.Return())
 }

core/syncx/timeoutlimit_test.go

@@ -29,5 +29,5 @@ func TestTimeoutLimit(t *testing.T) {
 	assert.Equal(t, ErrTimeout, limit.Borrow(time.Millisecond*100))
 	assert.Nil(t, limit.Return())
 	assert.Nil(t, limit.Return())
-	assert.Equal(t, ErrReturn, limit.Return())
+	assert.Equal(t, ErrLimitReturn, limit.Return())
 }

go.mod

@@ -58,6 +58,7 @@ require (
 	google.golang.org/grpc v1.29.1
 	google.golang.org/protobuf v1.25.0 // indirect
 	gopkg.in/cheggaaa/pb.v1 v1.0.28
+	gopkg.in/h2non/gock.v1 v1.0.15
 	gopkg.in/yaml.v2 v2.2.8
 	honnef.co/go/tools v0.0.1-2020.1.4 // indirect
 	sigs.k8s.io/yaml v1.2.0 // indirect

go.sum

@@ -141,6 +141,8 @@ github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0/go.mod h1:8NvIoxWQoOIhqOTXgf
 github.com/grpc-ecosystem/grpc-gateway v1.9.5/go.mod h1:vNeuVxBJEsws4ogUvrchl83t/GYV9WGTSLVdBhOQFDY=
 github.com/grpc-ecosystem/grpc-gateway v1.14.3 h1:OCJlWkOUoTnl0neNGlf4fUm3TmbEtguw7vR+nGtnDjY=
 github.com/grpc-ecosystem/grpc-gateway v1.14.3/go.mod h1:6CwZWGDSPRJidgKAtJVvND6soZe6fT7iteq8wDPdhb0=
+github.com/h2non/parth v0.0.0-20190131123155-b4df798d6542 h1:2VTzZjLZBgl62/EtslCrtky5vbi9dd7HrQPQIx6wqiw=
+github.com/h2non/parth v0.0.0-20190131123155-b4df798d6542/go.mod h1:Ow0tF8D4Kplbc8s8sSb3V2oUCygFHVp8gC3Dn6U4MNI=
 github.com/hpcloud/tail v1.0.0 h1:nfCOvKYfkgYP8hkirhJocXT2+zOD8yUNjXaWfTlyFKI=
 github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=
 github.com/iancoleman/strcase v0.0.0-20191112232945-16388991a334 h1:VHgatEHNcBFEB7inlalqfNqw65aNkM1lGX2yt3NmbS8=
@@ -205,6 +207,7 @@ github.com/modern-go/reflect2 v0.0.0-20180701023420-4b7aa43c6742/go.mod h1:bx2lN
 github.com/modern-go/reflect2 v1.0.1 h1:9f412s+6RmYXLWZSEzVVgPGK7C2PphHj5RJrvfx9AWI=
 github.com/modern-go/reflect2 v1.0.1/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3RllmbCylyMrvgv0=
 github.com/mwitkow/go-conntrack v0.0.0-20161129095857-cc309e4a2223/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U=
+github.com/nbio/st v0.0.0-20140626010706-e9e8d9816f32/go.mod h1:9wM+0iRr9ahx58uYLpLIr5fm8diHn0JbqRycJi6w0Ms=
 github.com/olekukonko/tablewriter v0.0.0-20170122224234-a0225b3f23b5/go.mod h1:vsDQFd/mU46D+Z4whnwzcISnGGzXWMclvtLoiIKAKIo=
 github.com/olekukonko/tablewriter v0.0.4 h1:vHD/YYe1Wolo78koG299f7V/VAS08c6IpCLn+Ejf/w8=
 github.com/olekukonko/tablewriter v0.0.4/go.mod h1:zq6QwlOf5SlnkVbMSr5EoBv3636FWnp+qbPhuoO21uA=
@@ -429,6 +432,8 @@ gopkg.in/cheggaaa/pb.v1 v1.0.28/go.mod h1:V/YB90LKu/1FcN3WVnfiiE5oMCibMjukxqG/qS
 gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI=
 gopkg.in/fsnotify.v1 v1.4.7 h1:xOHLXZwVvI9hhs+cLKq5+I5onOuwQLhQwiu63xxlHs4=
 gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys=
+gopkg.in/h2non/gock.v1 v1.0.15 h1:SzLqcIlb/fDfg7UvukMpNcWsu7sI5tWwL+KCATZqks0=
+gopkg.in/h2non/gock.v1 v1.0.15/go.mod h1:sX4zAkdYX1TRGJ2JY156cFspQn4yRWn6p9EMdODlynE=
 gopkg.in/resty.v1 v1.12.0/go.mod h1:mDo4pnntr5jdWRML875a/NmxYqAlA73dVijT2AXvQQo=
 gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 h1:uRGJdciOHaEIrze2W8Q3AKkepLTh2hOroT7a+7czfdQ=
 gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw=

readme.md

@@ -169,4 +169,4 @@ go get -u github.com/tal-tech/go-zero
 
 ## 9. 微信交流群
 
-<img src="doc/images/wechat.jpg" alt="wechat" width="300" />
+<img src="https://raw.githubusercontent.com/tal-tech/zero-doc/main/doc/images/wechat.jpg" alt="wechat" width="300" />

rest/handler/maxconnshandler.go

@@ -28,7 +28,7 @@ func MaxConns(n int) func(http.Handler) http.Handler {
 
 				next.ServeHTTP(w, r)
 			} else {
-				internal.Errorf(r, "Concurrent connections over %d, rejected with code %d",
+				internal.Errorf(r, "concurrent connections over %d, rejected with code %d",
 					n, http.StatusServiceUnavailable)
 				w.WriteHeader(http.StatusServiceUnavailable)
 			}

tools/goctl/api/parser/parser_test.go

@@ -0,0 +1,166 @@
+package parser
+
+import (
+	"io/ioutil"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+const testApiTemplate = `
+info(
+    title: doc title
+    desc: >
+    doc description first part,
+    doc description second part<
+    version: 1.0
+)
+
+type Request struct {
+  Name string ` + "`" + `path:"name,options=you|me"` + "`" + `
+}
+
+type Response struct {
+  Message string ` + "`" + `json:"message"` + "`" + `
+}
+
+service A-api {
+  @server(
+    handler: GreetHandler
+  )
+  get /greet/from/:name(Request) returns (Response)
+
+  @server(
+    handler: NoResponseHandler
+  )
+  get /greet/get(Request) returns
+}
+`
+
+const testMultiServiceTemplate = `
+info(
+    title: doc title
+    desc: >
+    doc description first part,
+    doc description second part<
+    version: 1.0
+)
+
+type Request struct {
+  Name string ` + "`" + `path:"name,options=you|me"` + "`" + ` 
+}
+
+type Response struct {
+  Message string ` + "`" + `json:"message"` + "`" + `
+}
+
+service A-api {
+  @server(
+    handler: GreetHandler
+  )
+  get /greet/from/:name(Request) returns (Response)
+}
+
+service A-api {
+  @server(
+    handler: NoResponseHandler
+  )
+  get /greet/get(Request) returns
+}
+`
+
+const apiNoInfo = `
+type Request struct {
+  Name string ` + "`" + `path:"name,options=you|me"` + "`" + `
+}
+
+type Response struct {
+  Message string ` + "`" + `json:"message"` + "`" + `
+}
+
+service A-api {
+  @server(
+    handler: GreetHandler
+  )
+  get /greet/from/:name(Request) returns (Response)
+}
+`
+
+const invalidApiFile = `
+type Request struct {
+  Name string ` + "`" + `path:"name,options=you|me"` + "`" + `
+}
+
+type Response struct {
+  Message string ` + "`" + `json:"message"` + "`" + `
+}
+
+service A-api
+  @server(
+    handler: GreetHandler
+  )
+  get /greet/from/:name(Request) returns (Response)
+}
+`
+
+func TestParser(t *testing.T) {
+	filename := "greet.api"
+	err := ioutil.WriteFile(filename, []byte(testApiTemplate), os.ModePerm)
+	assert.Nil(t, err)
+	defer os.Remove(filename)
+
+	parser, err := NewParser(filename)
+	assert.Nil(t, err)
+
+	api, err := parser.Parse()
+	assert.Nil(t, err)
+
+	assert.Equal(t, len(api.Types), 2)
+	assert.Equal(t, len(api.Service.Routes), 2)
+
+	assert.Equal(t, api.Service.Routes[0].Path, "/greet/from/:name")
+	assert.Equal(t, api.Service.Routes[1].Path, "/greet/get")
+
+	assert.Equal(t, api.Service.Routes[1].RequestType.Name, "Request")
+	assert.Equal(t, api.Service.Routes[1].ResponseType.Name, "")
+}
+
+func TestMultiService(t *testing.T) {
+	filename := "greet.api"
+	err := ioutil.WriteFile(filename, []byte(testMultiServiceTemplate), os.ModePerm)
+	assert.Nil(t, err)
+	defer os.Remove(filename)
+
+	parser, err := NewParser(filename)
+	assert.Nil(t, err)
+
+	api, err := parser.Parse()
+	assert.Nil(t, err)
+
+	assert.Equal(t, len(api.Service.Routes), 2)
+	assert.Equal(t, len(api.Service.Groups), 2)
+}
+
+func TestApiNoInfo(t *testing.T) {
+	filename := "greet.api"
+	err := ioutil.WriteFile(filename, []byte(apiNoInfo), os.ModePerm)
+	assert.Nil(t, err)
+	defer os.Remove(filename)
+
+	parser, err := NewParser(filename)
+	assert.Nil(t, err)
+
+	_, err = parser.Parse()
+	assert.Nil(t, err)
+}
+
+func TestInvalidApiFile(t *testing.T) {
+	filename := "greet.api"
+	err := ioutil.WriteFile(filename, []byte(invalidApiFile), os.ModePerm)
+	assert.Nil(t, err)
+	defer os.Remove(filename)
+
+	_, err = NewParser(filename)
+	assert.NotNil(t, err)
+}

zrpc/internal/auth/auth_test.go

@@ -0,0 +1,84 @@
+package auth
+
+import (
+	"context"
+	"testing"
+
+	"github.com/alicebob/miniredis"
+	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/stores/redis"
+	"google.golang.org/grpc/metadata"
+)
+
+func TestAuthenticator(t *testing.T) {
+	tests := []struct {
+		name     string
+		app      string
+		token    string
+		strict   bool
+		hasError bool
+	}{
+		{
+			name:     "strict=false",
+			strict:   false,
+			hasError: false,
+		},
+		{
+			name:     "strict=true",
+			strict:   true,
+			hasError: true,
+		},
+		{
+			name:     "strict=true,with token",
+			app:      "foo",
+			token:    "bar",
+			strict:   true,
+			hasError: false,
+		},
+		{
+			name:     "strict=true,with error token",
+			app:      "foo",
+			token:    "error",
+			strict:   true,
+			hasError: true,
+		},
+	}
+
+	r := miniredis.NewMiniRedis()
+	assert.Nil(t, r.Start())
+	defer r.Close()
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			store := redis.NewRedis(r.Addr(), redis.NodeType)
+			if len(test.app) > 0 {
+				assert.Nil(t, store.Hset("apps", test.app, test.token))
+				defer store.Hdel("apps", test.app)
+			}
+
+			authenticator, err := NewAuthenticator(store, "apps", test.strict)
+			assert.Nil(t, err)
+			assert.NotNil(t, authenticator.Authenticate(context.Background()))
+			md := metadata.New(map[string]string{})
+			ctx := metadata.NewIncomingContext(context.Background(), md)
+			assert.NotNil(t, authenticator.Authenticate(ctx))
+			md = metadata.New(map[string]string{
+				"app":   "",
+				"token": "",
+			})
+			ctx = metadata.NewIncomingContext(context.Background(), md)
+			assert.NotNil(t, authenticator.Authenticate(ctx))
+			md = metadata.New(map[string]string{
+				"app":   "foo",
+				"token": "bar",
+			})
+			ctx = metadata.NewIncomingContext(context.Background(), md)
+			err = authenticator.Authenticate(ctx)
+			if test.hasError {
+				assert.NotNil(t, err)
+			} else {
+				assert.Nil(t, err)
+			}
+		})
+	}
+}