Skip to content

Commit b41b1b0

Browse files
me-csme-cs
authored
update:github.com/mongodb/mongo-go-driver v2.0 Migration (#4687)
1 parent f36e5fe commit b41b1b0

22 files changed

+2519
-1178
lines changed

‎core/stores/mon/bulkinserter.go‎

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
//go:generate mockgen -package mon -destination collection_inserter_mock.go -source bulkinserter.go collectionInserter
12
package mon
23

34
import (
@@ -6,7 +7,8 @@ import (
67

78
"github.com/zeromicro/go-zero/core/executors"
89
"github.com/zeromicro/go-zero/core/logx"
9-
"go.mongodb.org/mongo-driver/mongo"
10+
"go.mongodb.org/mongo-driver/v2/mongo"
11+
"go.mongodb.org/mongo-driver/v2/mongo/options"
1012
)
1113

1214
const (
@@ -27,10 +29,7 @@ type (
2729

2830
// NewBulkInserter returns a BulkInserter.
2931
func NewBulkInserter(coll Collection, interval ...time.Duration) (*BulkInserter, error) {
30-
cloneColl, err := coll.Clone()
31-
if err != nil {
32-
return nil, err
33-
}
32+
cloneColl := coll.Clone()
3433

3534
inserter := &dbInserter{
3635
collection: cloneColl,
@@ -64,8 +63,16 @@ func (bi *BulkInserter) SetResultHandler(handler ResultHandler) {
6463
})
6564
}
6665

66+
type collectionInserter interface {
67+
InsertMany(
68+
ctx context.Context,
69+
documents interface{},
70+
opts ...options.Lister[options.InsertManyOptions],
71+
) (*mongo.InsertManyResult, error)
72+
}
73+
6774
type dbInserter struct {
68-
collection *mongo.Collection
75+
collection collectionInserter
6976
documents []any
7077
resultHandler ResultHandler
7178
}

‎core/stores/mon/bulkinserter_test.go‎

Lines changed: 121 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,26 +1,131 @@
11
package mon
22

33
import (
4+
"errors"
45
"testing"
6+
"time"
57

68
"github.com/stretchr/testify/assert"
7-
"go.mongodb.org/mongo-driver/bson"
8-
"go.mongodb.org/mongo-driver/mongo"
9-
"go.mongodb.org/mongo-driver/mongo/integration/mtest"
9+
"go.mongodb.org/mongo-driver/v2/mongo"
10+
"go.uber.org/mock/gomock"
1011
)
1112

12-
func TestBulkInserter(t *testing.T) {
13-
mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
14-
mt.Run("test", func(mt *mtest.T) {
15-
mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "ok", Value: 1}}...))
16-
bulk, err := NewBulkInserter(createModel(mt).Collection)
17-
assert.Equal(t, err, nil)
18-
bulk.SetResultHandler(func(result *mongo.InsertManyResult, err error) {
19-
assert.Nil(t, err)
20-
assert.Equal(t, 2, len(result.InsertedIDs))
21-
})
22-
bulk.Insert(bson.D{{Key: "foo", Value: "bar"}})
23-
bulk.Insert(bson.D{{Key: "foo", Value: "baz"}})
24-
bulk.Flush()
13+
func TestBulkInserter_InsertAndFlush(t *testing.T) {
14+
ctrl := gomock.NewController(t)
15+
defer ctrl.Finish()
16+
mockCollection := NewMockCollection(ctrl)
17+
mockCollection.EXPECT().Clone().Return(&mongo.Collection{})
18+
bulkInserter, err := NewBulkInserter(mockCollection, time.Second)
19+
assert.NoError(t, err)
20+
bulkInserter.SetResultHandler(func(result *mongo.InsertManyResult, err error) {
21+
assert.Nil(t, err)
22+
assert.Equal(t, 2, len(result.InsertedIDs))
2523
})
24+
doc := map[string]interface{}{"name": "test"}
25+
bulkInserter.Insert(doc)
26+
bulkInserter.Flush()
27+
}
28+
29+
func TestBulkInserter_SetResultHandler(t *testing.T) {
30+
ctrl := gomock.NewController(t)
31+
defer ctrl.Finish()
32+
mockCollection := NewMockCollection(ctrl)
33+
mockCollection.EXPECT().Clone().Return(nil)
34+
bulkInserter, err := NewBulkInserter(mockCollection)
35+
assert.NoError(t, err)
36+
mockHandler := func(result *mongo.InsertManyResult, err error) {}
37+
bulkInserter.SetResultHandler(mockHandler)
38+
}
39+
40+
func TestDbInserter_RemoveAll(t *testing.T) {
41+
inserter := &dbInserter{}
42+
inserter.documents = []interface{}{}
43+
docs := inserter.RemoveAll()
44+
assert.NotNil(t, docs)
45+
assert.Empty(t, inserter.documents)
46+
}
47+
48+
func Test_dbInserter_Execute(t *testing.T) {
49+
type fields struct {
50+
collection collectionInserter
51+
documents []any
52+
resultHandler ResultHandler
53+
}
54+
ctrl := gomock.NewController(t)
55+
defer ctrl.Finish()
56+
mockCollection := NewMockcollectionInserter(ctrl)
57+
type args struct {
58+
objs any
59+
}
60+
tests := []struct {
61+
name string
62+
fields fields
63+
args args
64+
mock func()
65+
}{
66+
{
67+
name: "empty doc",
68+
fields: fields{
69+
collection: nil,
70+
documents: nil,
71+
resultHandler: nil,
72+
},
73+
args: args{
74+
objs: make([]any, 0),
75+
},
76+
mock: func() {},
77+
},
78+
{
79+
name: "result handler",
80+
fields: fields{
81+
collection: mockCollection,
82+
resultHandler: func(result *mongo.InsertManyResult, err error) {
83+
assert.NotNil(t, err)
84+
},
85+
},
86+
args: args{
87+
objs: make([]any, 1),
88+
},
89+
mock: func() {
90+
mockCollection.EXPECT().InsertMany(gomock.Any(), gomock.Any()).Return(&mongo.InsertManyResult{}, errors.New("error"))
91+
},
92+
},
93+
{
94+
name: "normal error handler",
95+
fields: fields{
96+
collection: mockCollection,
97+
resultHandler: nil,
98+
},
99+
args: args{
100+
objs: make([]any, 1),
101+
},
102+
mock: func() {
103+
mockCollection.EXPECT().InsertMany(gomock.Any(), gomock.Any()).Return(&mongo.InsertManyResult{}, errors.New("error"))
104+
},
105+
},
106+
{
107+
name: "no error",
108+
fields: fields{
109+
collection: mockCollection,
110+
resultHandler: nil,
111+
},
112+
args: args{
113+
objs: make([]any, 1),
114+
},
115+
mock: func() {
116+
mockCollection.EXPECT().InsertMany(gomock.Any(), gomock.Any()).Return(&mongo.InsertManyResult{}, nil)
117+
},
118+
},
119+
}
120+
for _, tt := range tests {
121+
t.Run(tt.name, func(t *testing.T) {
122+
tt.mock()
123+
in := &dbInserter{
124+
collection: tt.fields.collection,
125+
documents: tt.fields.documents,
126+
resultHandler: tt.fields.resultHandler,
127+
}
128+
in.Execute(tt.args.objs)
129+
})
130+
}
26131
}

‎core/stores/mon/clientmanager.go‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,8 @@ import (
55
"io"
66

77
"github.com/zeromicro/go-zero/core/syncx"
8-
"go.mongodb.org/mongo-driver/mongo"
9-
mopt "go.mongodb.org/mongo-driver/mongo/options"
8+
"go.mongodb.org/mongo-driver/v2/mongo"
9+
"go.mongodb.org/mongo-driver/v2/mongo/options"
1010
)
1111

1212
var clientManager = syncx.NewResourceManager()
@@ -29,13 +29,13 @@ func Inject(key string, client *mongo.Client) {
2929

3030
func getClient(url string, opts ...Option) (*mongo.Client, error) {
3131
val, err := clientManager.GetResource(url, func() (io.Closer, error) {
32-
o := mopt.Client().ApplyURI(url)
32+
o := options.Client().ApplyURI(url)
3333
opts = append([]Option{defaultTimeoutOption()}, opts...)
3434
for _, opt := range opts {
3535
opt(o)
3636
}
3737

38-
cli, err := mongo.Connect(context.Background(), o)
38+
cli, err := mongo.Connect(o)
3939
if err != nil {
4040
return nil, err
4141
}

‎core/stores/mon/clientmanager_test.go‎

Lines changed: 6 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -4,19 +4,13 @@ import (
44
"testing"
55

66
"github.com/stretchr/testify/assert"
7-
"go.mongodb.org/mongo-driver/mongo/integration/mtest"
7+
"go.mongodb.org/mongo-driver/v2/mongo"
88
)
99

10-
func init() {
11-
_ = mtest.Setup()
12-
}
13-
1410
func TestClientManger_getClient(t *testing.T) {
15-
mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
16-
mt.Run("test", func(mt *mtest.T) {
17-
Inject(mtest.ClusterURI(), mt.Client)
18-
cli, err := getClient(mtest.ClusterURI())
19-
assert.Nil(t, err)
20-
assert.Equal(t, mt.Client, cli)
21-
})
11+
c := &mongo.Client{}
12+
Inject("foo", c)
13+
cli, err := getClient("foo")
14+
assert.Nil(t, err)
15+
assert.Equal(t, c, cli)
2216
}

0 commit comments

Comments
 (0)