diff --git a/core/stores/mongo/collection.go b/core/stores/mongo/collection.go index ee56efd5..aee8a972 100644 --- a/core/stores/mongo/collection.go +++ b/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)) } } } diff --git a/core/stores/mongo/collection_test.go b/core/stores/mongo/collection_test.go index 31cb10ec..2b6c7825 100644 --- a/core/stores/mongo/collection_test.go +++ b/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 +} diff --git a/core/stores/mongo/internal/collection.go b/core/stores/mongo/internal/collection.go new file mode 100644 index 00000000..c6ca237d --- /dev/null +++ b/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) +} diff --git a/core/stores/mongo/internal/collection_mock.go b/core/stores/mongo/internal/collection_mock.go new file mode 100644 index 00000000..5f6b7b40 --- /dev/null +++ b/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) +}