You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
go-zero/core/stores/mon/collection_test.go

614 lines
18 KiB
Go

package mon
import (
"context"
"errors"
"testing"
"github.com/stretchr/testify/assert"
"github.com/zeromicro/go-zero/core/breaker"
"github.com/zeromicro/go-zero/core/logx/logtest"
"github.com/zeromicro/go-zero/core/stringx"
"github.com/zeromicro/go-zero/core/timex"
"go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/bson/primitive"
"go.mongodb.org/mongo-driver/mongo"
"go.mongodb.org/mongo-driver/mongo/integration/mtest"
mopt "go.mongodb.org/mongo-driver/mongo/options"
)
var errDummy = errors.New("dummy")
func TestKeepPromise_accept(t *testing.T) {
p := new(mockPromise)
kp := keepablePromise{
promise: p,
log: func(error) {},
}
assert.Nil(t, kp.accept(nil))
assert.Equal(t, ErrNotFound, kp.accept(ErrNotFound))
}
func TestKeepPromise_keep(t *testing.T) {
tests := []struct {
err error
accepted bool
reason string
}{
{
err: nil,
accepted: true,
reason: "",
},
{
err: ErrNotFound,
accepted: true,
reason: "",
},
{
err: errors.New("any"),
accepted: false,
reason: "any",
},
}
for _, test := range tests {
t.Run(stringx.RandId(), func(t *testing.T) {
p := new(mockPromise)
kp := keepablePromise{
promise: p,
log: func(error) {},
}
assert.Equal(t, test.err, kp.keep(test.err))
assert.Equal(t, test.accepted, p.accepted)
assert.Equal(t, test.reason, p.reason)
})
}
}
func TestNewCollection(t *testing.T) {
mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
mt.Run("test", func(mt *mtest.T) {
coll := mt.Coll
assert.NotNil(t, coll)
col := newCollection(coll, breaker.GetBreaker("localhost"))
assert.Equal(t, t.Name()+"/test", col.(*decoratedCollection).name)
})
}
func TestCollection_Aggregate(t *testing.T) {
mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
mt.Run("test", func(mt *mtest.T) {
coll := mt.Coll
assert.NotNil(t, coll)
col := newCollection(coll, breaker.GetBreaker("localhost"))
ns := mt.Coll.Database().Name() + "." + mt.Coll.Name()
aggRes := mtest.CreateCursorResponse(1, ns, mtest.FirstBatch)
mt.AddMockResponses(aggRes)
assert.Equal(t, t.Name()+"/test", col.(*decoratedCollection).name)
cursor, err := col.Aggregate(context.Background(), mongo.Pipeline{}, mopt.Aggregate())
assert.Nil(t, err)
cursor.Close(context.Background())
})
}
func TestCollection_BulkWrite(t *testing.T) {
mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
mt.Run("test", func(mt *mtest.T) {
c := decoratedCollection{
Collection: mt.Coll,
brk: breaker.NewBreaker(),
}
mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "ok", Value: 1}}...))
res, err := c.BulkWrite(context.Background(), []mongo.WriteModel{
mongo.NewInsertOneModel().SetDocument(bson.D{{Key: "foo", Value: 1}}),
})
assert.Nil(t, err)
assert.NotNil(t, res)
c.brk = new(dropBreaker)
_, err = c.BulkWrite(context.Background(), []mongo.WriteModel{
mongo.NewInsertOneModel().SetDocument(bson.D{{Key: "foo", Value: 1}}),
})
assert.Equal(t, errDummy, err)
})
}
func TestCollection_CountDocuments(t *testing.T) {
mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
mt.Run("test", func(mt *mtest.T) {
c := decoratedCollection{
Collection: mt.Coll,
brk: breaker.NewBreaker(),
}
mt.AddMockResponses(mtest.CreateCursorResponse(
1,
"DBName.CollectionName",
mtest.FirstBatch,
bson.D{
{Key: "n", Value: 1},
}))
res, err := c.CountDocuments(context.Background(), bson.D{})
assert.Nil(t, err)
assert.Equal(t, int64(1), res)
c.brk = new(dropBreaker)
_, err = c.CountDocuments(context.Background(), bson.D{{Key: "foo", Value: 1}})
assert.Equal(t, errDummy, err)
})
}
func TestDecoratedCollection_DeleteMany(t *testing.T) {
mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
mt.Run("test", func(mt *mtest.T) {
c := decoratedCollection{
Collection: mt.Coll,
brk: breaker.NewBreaker(),
}
mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
res, err := c.DeleteMany(context.Background(), bson.D{})
assert.Nil(t, err)
assert.Equal(t, int64(1), res.DeletedCount)
c.brk = new(dropBreaker)
_, err = c.DeleteMany(context.Background(), bson.D{{Key: "foo", Value: 1}})
assert.Equal(t, errDummy, err)
})
}
func TestCollection_Distinct(t *testing.T) {
mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
mt.Run("test", func(mt *mtest.T) {
c := decoratedCollection{
Collection: mt.Coll,
brk: breaker.NewBreaker(),
}
mt.AddMockResponses(bson.D{{Key: "ok", Value: 1}, {Key: "values", Value: []int{1}}})
resp, err := c.Distinct(context.Background(), "foo", bson.D{})
assert.Nil(t, err)
assert.Equal(t, 1, len(resp))
c.brk = new(dropBreaker)
_, err = c.Distinct(context.Background(), "foo", bson.D{{Key: "foo", Value: 1}})
assert.Equal(t, errDummy, err)
})
}
func TestCollection_EstimatedDocumentCount(t *testing.T) {
mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
mt.Run("test", func(mt *mtest.T) {
c := decoratedCollection{
Collection: mt.Coll,
brk: breaker.NewBreaker(),
}
mt.AddMockResponses(bson.D{{Key: "ok", Value: 1}, {Key: "n", Value: 1}})
res, err := c.EstimatedDocumentCount(context.Background())
assert.Nil(t, err)
assert.Equal(t, int64(1), res)
c.brk = new(dropBreaker)
_, err = c.EstimatedDocumentCount(context.Background())
assert.Equal(t, errDummy, err)
})
}
func TestCollection_Find(t *testing.T) {
mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
mt.Run("test", func(mt *mtest.T) {
c := decoratedCollection{
Collection: mt.Coll,
brk: breaker.NewBreaker(),
}
find := mtest.CreateCursorResponse(
1,
"DBName.CollectionName",
mtest.FirstBatch,
bson.D{
{Key: "name", Value: "John"},
})
getMore := mtest.CreateCursorResponse(
1,
"DBName.CollectionName",
mtest.NextBatch,
bson.D{
{Key: "name", Value: "Mary"},
})
killCursors := mtest.CreateCursorResponse(
0,
"DBName.CollectionName",
mtest.NextBatch)
mt.AddMockResponses(find, getMore, killCursors)
filter := bson.D{{Key: "x", Value: 1}}
cursor, err := c.Find(context.Background(), filter, mopt.Find())
assert.Nil(t, err)
defer cursor.Close(context.Background())
var val []struct {
ID primitive.ObjectID `bson:"_id"`
Name string `bson:"name"`
}
assert.Nil(t, cursor.All(context.Background(), &val))
assert.Equal(t, 2, len(val))
assert.Equal(t, "John", val[0].Name)
assert.Equal(t, "Mary", val[1].Name)
c.brk = new(dropBreaker)
_, err = c.Find(context.Background(), filter, mopt.Find())
assert.Equal(t, errDummy, err)
})
}
func TestCollection_FindOne(t *testing.T) {
mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
mt.Run("test", func(mt *mtest.T) {
c := decoratedCollection{
Collection: mt.Coll,
brk: breaker.NewBreaker(),
}
find := mtest.CreateCursorResponse(
1,
"DBName.CollectionName",
mtest.FirstBatch,
bson.D{
{Key: "name", Value: "John"},
})
getMore := mtest.CreateCursorResponse(
1,
"DBName.CollectionName",
mtest.NextBatch,
bson.D{
{Key: "name", Value: "Mary"},
})
killCursors := mtest.CreateCursorResponse(
0,
"DBName.CollectionName",
mtest.NextBatch)
mt.AddMockResponses(find, getMore, killCursors)
filter := bson.D{{Key: "x", Value: 1}}
resp, err := c.FindOne(context.Background(), filter)
assert.Nil(t, err)
var val struct {
ID primitive.ObjectID `bson:"_id"`
Name string `bson:"name"`
}
assert.Nil(t, resp.Decode(&val))
assert.Equal(t, "John", val.Name)
c.brk = new(dropBreaker)
_, err = c.FindOne(context.Background(), filter)
assert.Equal(t, errDummy, err)
})
}
func TestCollection_FindOneAndDelete(t *testing.T) {
mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
mt.Run("test", func(mt *mtest.T) {
c := decoratedCollection{
Collection: mt.Coll,
brk: breaker.NewBreaker(),
}
filter := bson.D{}
mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{}...))
_, err := c.FindOneAndDelete(context.Background(), filter, mopt.FindOneAndDelete())
assert.Equal(t, mongo.ErrNoDocuments, err)
mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{
{Key: "value", Value: bson.D{{Key: "name", Value: "John"}}},
}...))
resp, err := c.FindOneAndDelete(context.Background(), filter, mopt.FindOneAndDelete())
assert.Nil(t, err)
var val struct {
Name string `bson:"name"`
}
assert.Nil(t, resp.Decode(&val))
assert.Equal(t, "John", val.Name)
c.brk = new(dropBreaker)
_, err = c.FindOneAndDelete(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
assert.Equal(t, errDummy, err)
})
}
func TestCollection_FindOneAndReplace(t *testing.T) {
mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
mt.Run("test", func(mt *mtest.T) {
c := decoratedCollection{
Collection: mt.Coll,
brk: breaker.NewBreaker(),
}
mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{}...))
filter := bson.D{{Key: "x", Value: 1}}
replacement := bson.D{{Key: "x", Value: 2}}
opts := mopt.FindOneAndReplace().SetUpsert(true)
_, err := c.FindOneAndReplace(context.Background(), filter, replacement, opts)
assert.Equal(t, mongo.ErrNoDocuments, err)
mt.AddMockResponses(bson.D{{Key: "ok", Value: 1}, {Key: "value", Value: bson.D{
{Key: "name", Value: "John"},
}}})
resp, err := c.FindOneAndReplace(context.Background(), filter, replacement, opts)
assert.Nil(t, err)
var val struct {
Name string `bson:"name"`
}
assert.Nil(t, resp.Decode(&val))
assert.Equal(t, "John", val.Name)
c.brk = new(dropBreaker)
_, err = c.FindOneAndReplace(context.Background(), filter, replacement, opts)
assert.Equal(t, errDummy, err)
})
}
func TestCollection_FindOneAndUpdate(t *testing.T) {
mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
mt.Run("test", func(mt *mtest.T) {
c := decoratedCollection{
Collection: mt.Coll,
brk: breaker.NewBreaker(),
}
mt.AddMockResponses(bson.D{{Key: "ok", Value: 1}})
filter := bson.D{{Key: "x", Value: 1}}
update := bson.D{{Key: "$x", Value: 2}}
opts := mopt.FindOneAndUpdate().SetUpsert(true)
_, err := c.FindOneAndUpdate(context.Background(), filter, update, opts)
assert.Equal(t, mongo.ErrNoDocuments, err)
mt.AddMockResponses(bson.D{{Key: "ok", Value: 1}, {Key: "value", Value: bson.D{
{Key: "name", Value: "John"},
}}})
resp, err := c.FindOneAndUpdate(context.Background(), filter, update, opts)
assert.Nil(t, err)
var val struct {
Name string `bson:"name"`
}
assert.Nil(t, resp.Decode(&val))
assert.Equal(t, "John", val.Name)
c.brk = new(dropBreaker)
_, err = c.FindOneAndUpdate(context.Background(), filter, update, opts)
assert.Equal(t, errDummy, err)
})
}
func TestCollection_InsertOne(t *testing.T) {
mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
mt.Run("test", func(mt *mtest.T) {
c := decoratedCollection{
Collection: mt.Coll,
brk: breaker.NewBreaker(),
}
mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "ok", Value: 1}}...))
res, err := c.InsertOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
assert.Nil(t, err)
assert.NotNil(t, res)
c.brk = new(dropBreaker)
_, err = c.InsertOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
assert.Equal(t, errDummy, err)
})
}
func TestCollection_InsertMany(t *testing.T) {
mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
mt.Run("test", func(mt *mtest.T) {
c := decoratedCollection{
Collection: mt.Coll,
brk: breaker.NewBreaker(),
}
mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "ok", Value: 1}}...))
res, err := c.InsertMany(context.Background(), []any{
bson.D{{Key: "foo", Value: "bar"}},
bson.D{{Key: "foo", Value: "baz"}},
})
assert.Nil(t, err)
assert.NotNil(t, res)
assert.Equal(t, 2, len(res.InsertedIDs))
c.brk = new(dropBreaker)
_, err = c.InsertMany(context.Background(), []any{bson.D{{Key: "foo", Value: "bar"}}})
assert.Equal(t, errDummy, err)
})
}
func TestCollection_DeleteOne(t *testing.T) {
mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
mt.Run("test", func(mt *mtest.T) {
c := decoratedCollection{
Collection: mt.Coll,
brk: breaker.NewBreaker(),
}
mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
res, err := c.DeleteOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
assert.Nil(t, err)
assert.Equal(t, int64(1), res.DeletedCount)
c.brk = new(dropBreaker)
_, err = c.DeleteOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
assert.Equal(t, errDummy, err)
})
}
func TestCollection_DeleteMany(t *testing.T) {
mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
mt.Run("test", func(mt *mtest.T) {
c := decoratedCollection{
Collection: mt.Coll,
brk: breaker.NewBreaker(),
}
mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
res, err := c.DeleteMany(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
assert.Nil(t, err)
assert.Equal(t, int64(1), res.DeletedCount)
c.brk = new(dropBreaker)
_, err = c.DeleteMany(context.Background(), bson.D{{Key: "foo", Value: "bar"}})
assert.Equal(t, errDummy, err)
})
}
func TestCollection_ReplaceOne(t *testing.T) {
mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
mt.Run("test", func(mt *mtest.T) {
c := decoratedCollection{
Collection: mt.Coll,
brk: breaker.NewBreaker(),
}
mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
res, err := c.ReplaceOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
bson.D{{Key: "foo", Value: "baz"}},
)
assert.Nil(t, err)
assert.Equal(t, int64(1), res.MatchedCount)
c.brk = new(dropBreaker)
_, err = c.ReplaceOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
bson.D{{Key: "foo", Value: "baz"}})
assert.Equal(t, errDummy, err)
})
}
func TestCollection_UpdateOne(t *testing.T) {
mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
mt.Run("test", func(mt *mtest.T) {
c := decoratedCollection{
Collection: mt.Coll,
brk: breaker.NewBreaker(),
}
mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
resp, err := c.UpdateOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
assert.Nil(t, err)
assert.Equal(t, int64(1), resp.MatchedCount)
c.brk = new(dropBreaker)
_, err = c.UpdateOne(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
assert.Equal(t, errDummy, err)
})
}
func TestCollection_UpdateByID(t *testing.T) {
mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
mt.Run("test", func(mt *mtest.T) {
c := decoratedCollection{
Collection: mt.Coll,
brk: breaker.NewBreaker(),
}
mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
resp, err := c.UpdateByID(context.Background(), primitive.NewObjectID(),
bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
assert.Nil(t, err)
assert.Equal(t, int64(1), resp.MatchedCount)
c.brk = new(dropBreaker)
_, err = c.UpdateByID(context.Background(), primitive.NewObjectID(),
bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
assert.Equal(t, errDummy, err)
})
}
func TestCollection_UpdateMany(t *testing.T) {
mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
mt.Run("test", func(mt *mtest.T) {
c := decoratedCollection{
Collection: mt.Coll,
brk: breaker.NewBreaker(),
}
mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...))
resp, err := c.UpdateMany(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
assert.Nil(t, err)
assert.Equal(t, int64(1), resp.MatchedCount)
c.brk = new(dropBreaker)
_, err = c.UpdateMany(context.Background(), bson.D{{Key: "foo", Value: "bar"}},
bson.D{{Key: "$set", Value: bson.D{{Key: "baz", Value: "qux"}}}})
assert.Equal(t, errDummy, err)
})
}
func TestDecoratedCollection_LogDuration(t *testing.T) {
mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
c := decoratedCollection{
Collection: mt.Coll,
brk: breaker.NewBreaker(),
}
buf := logtest.NewCollector(t)
buf.Reset()
c.logDuration(context.Background(), "foo", timex.Now(), nil, "bar")
assert.Contains(t, buf.String(), "foo")
assert.Contains(t, buf.String(), "bar")
buf.Reset()
c.logDuration(context.Background(), "foo", timex.Now(), errors.New("bar"), make(chan int))
assert.Contains(t, buf.String(), "foo")
assert.Contains(t, buf.String(), "bar")
buf.Reset()
c.logDuration(context.Background(), "foo", timex.Now(), nil, make(chan int))
assert.Contains(t, buf.String(), "foo")
buf.Reset()
c.logDuration(context.Background(), "foo", timex.Now()-slowThreshold.Load()*2,
nil, make(chan int))
assert.Contains(t, buf.String(), "foo")
assert.Contains(t, buf.String(), "slowcall")
buf.Reset()
c.logDuration(context.Background(), "foo", timex.Now()-slowThreshold.Load()*2,
errors.New("bar"), make(chan int))
assert.Contains(t, buf.String(), "foo")
assert.Contains(t, buf.String(), "bar")
buf.Reset()
c.logDuration(context.Background(), "foo", timex.Now()-slowThreshold.Load()*2,
errors.New("bar"))
assert.Contains(t, buf.String(), "foo")
buf.Reset()
c.logDuration(context.Background(), "foo", timex.Now()-slowThreshold.Load()*2, nil)
assert.Contains(t, buf.String(), "foo")
assert.Contains(t, buf.String(), "slowcall")
}
type mockPromise struct {
accepted bool
reason string
}
func (p *mockPromise) Accept() {
p.accepted = true
}
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(_ func() error) error {
return nil
}
func (d *dropBreaker) DoWithAcceptable(_ func() error, _ breaker.Acceptable) error {
return errDummy
}
func (d *dropBreaker) DoWithFallback(_ func() error, _ breaker.Fallback) error {
return nil
}
func (d *dropBreaker) DoWithFallbackAcceptable(_ func() error, _ breaker.Fallback,
_ breaker.Acceptable) error {
return nil
}