milvus-io
/
milvus
mirror of https://github.com/milvus-io/milvus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Remove dup mock_kv.go and mock_txn_kv.go (#20267) 

Replace them with mockery generated mocks in kv/mocks

See also: #19289

Signed-off-by: yangxuan <xuan.yang@zilliz.com>

Signed-off-by: yangxuan <xuan.yang@zilliz.com>
pull/20292/head
XuanYang-cn 2022-11-03 14:41:35 +08:00 committed by GitHub
parent cb2591d1fe
commit 22fb1d3b93
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
14 changed files with 188 additions and 477 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
internal/datacoord/garbage_collector_test.go
View File

@ -45,8 +45,7 @@ func Test_garbageCollector_basic(t *testing.T) {
cli, _, _, _, _, err := initUtOSSEnv(bucketName, rootPath, 0)
require.NoError(t, err)
mockAllocator := newMockAllocator()
meta, err := newMemoryMeta(mockAllocator)
meta, err := newMemoryMeta()
assert.Nil(t, err)
etcdCli, err := etcd.GetEtcdClient(&Params.EtcdCfg)
@ -108,8 +107,7 @@ func Test_garbageCollector_scan(t *testing.T) {
cli, inserts, stats, delta, others, err := initUtOSSEnv(bucketName, rootPath, 4)
require.NoError(t, err)
mockAllocator := newMockAllocator()
meta, err := newMemoryMeta(mockAllocator)
meta, err := newMemoryMeta()
assert.Nil(t, err)
etcdCli, err := etcd.GetEtcdClient(&Params.EtcdCfg)

5
internal/datacoord/meta_test.go
View File

@ -243,7 +243,7 @@ func TestMeta_Basic(t *testing.T) {
ctx := context.Background()
mockAllocator := newMockAllocator()
meta, err := newMemoryMeta(mockAllocator)
meta, err := newMemoryMeta()
assert.Nil(t, err)
testSchema := newTestSchema()
@ -489,8 +489,7 @@ func TestMeta_Basic(t *testing.T) {
}
func TestGetUnFlushedSegments(t *testing.T) {
mockAllocator := newMockAllocator()
meta, err := newMemoryMeta(mockAllocator)
meta, err := newMemoryMeta()
assert.Nil(t, err)
s1 := &datapb.SegmentInfo{
ID: 0,

2
internal/datacoord/mock_test.go
View File

@ -38,7 +38,7 @@ import (
"github.com/milvus-io/milvus/internal/util/typeutil"
)
func newMemoryMeta(allocator allocator) (*meta, error) {
func newMemoryMeta() (*meta, error) {
memoryKV := memkv.NewMemoryKV()
return newMeta(context.TODO(), memoryKV, "")
}

26
internal/datacoord/segment_manager_test.go
View File

@ -35,7 +35,7 @@ func TestManagerOptions(t *testing.T) {
// ctx := context.Background()
Params.Init()
mockAllocator := newMockAllocator()
meta, err := newMemoryMeta(mockAllocator)
meta, err := newMemoryMeta()
assert.Nil(t, err)
segmentManager := newSegmentManager(meta, mockAllocator, nil)
@ -96,7 +96,7 @@ func TestAllocSegment(t *testing.T) {
ctx := context.Background()
Params.Init()
mockAllocator := newMockAllocator()
meta, err := newMemoryMeta(mockAllocator)
meta, err := newMemoryMeta()
assert.Nil(t, err)
segmentManager := newSegmentManager(meta, mockAllocator, nil)
@ -137,7 +137,7 @@ func TestAllocSegmentForImport(t *testing.T) {
ctx := context.Background()
Params.Init()
mockAllocator := newMockAllocator()
meta, err := newMemoryMeta(mockAllocator)
meta, err := newMemoryMeta()
assert.Nil(t, err)
ms := newMockRootCoordService()
segmentManager := newSegmentManager(meta, mockAllocator, ms)
@ -185,7 +185,7 @@ func TestLoadSegmentsFromMeta(t *testing.T) {
ctx := context.Background()
Params.Init()
mockAllocator := newMockAllocator()
meta, err := newMemoryMeta(mockAllocator)
meta, err := newMemoryMeta()
assert.Nil(t, err)
schema := newTestSchema()
@ -235,7 +235,7 @@ func TestLoadSegmentsFromMeta(t *testing.T) {
func TestSaveSegmentsToMeta(t *testing.T) {
Params.Init()
mockAllocator := newMockAllocator()
meta, err := newMemoryMeta(mockAllocator)
meta, err := newMemoryMeta()
assert.Nil(t, err)
schema := newTestSchema()
@ -257,7 +257,7 @@ func TestSaveSegmentsToMeta(t *testing.T) {
func TestSaveSegmentsToMetaWithSpecificSegments(t *testing.T) {
Params.Init()
mockAllocator := newMockAllocator()
meta, err := newMemoryMeta(mockAllocator)
meta, err := newMemoryMeta()
assert.Nil(t, err)
schema := newTestSchema()
@ -279,7 +279,7 @@ func TestSaveSegmentsToMetaWithSpecificSegments(t *testing.T) {
func TestDropSegment(t *testing.T) {
Params.Init()
mockAllocator := newMockAllocator()
meta, err := newMemoryMeta(mockAllocator)
meta, err := newMemoryMeta()
assert.Nil(t, err)
schema := newTestSchema()
@ -302,7 +302,7 @@ func TestDropSegment(t *testing.T) {
func TestAllocRowsLargerThanOneSegment(t *testing.T) {
Params.Init()
mockAllocator := newMockAllocator()
meta, err := newMemoryMeta(mockAllocator)
meta, err := newMemoryMeta()
assert.Nil(t, err)
schema := newTestSchema()
@ -324,7 +324,7 @@ func TestAllocRowsLargerThanOneSegment(t *testing.T) {
func TestExpireAllocation(t *testing.T) {
Params.Init()
mockAllocator := newMockAllocator()
meta, err := newMemoryMeta(mockAllocator)
meta, err := newMemoryMeta()
assert.Nil(t, err)
schema := newTestSchema()
@ -367,7 +367,7 @@ func TestGetFlushableSegments(t *testing.T) {
t.Run("get flushable segments between small interval", func(t *testing.T) {
Params.Init()
mockAllocator := newMockAllocator()
meta, err := newMemoryMeta(mockAllocator)
meta, err := newMemoryMeta()
assert.Nil(t, err)
schema := newTestSchema()
@ -413,7 +413,7 @@ func TestTryToSealSegment(t *testing.T) {
t.Run("normal seal with segment policies", func(t *testing.T) {
Params.Init()
mockAllocator := newMockAllocator()
meta, err := newMemoryMeta(mockAllocator)
meta, err := newMemoryMeta()
assert.Nil(t, err)
schema := newTestSchema()
@ -438,7 +438,7 @@ func TestTryToSealSegment(t *testing.T) {
t.Run("normal seal with channel seal policies", func(t *testing.T) {
Params.Init()
mockAllocator := newMockAllocator()
meta, err := newMemoryMeta(mockAllocator)
meta, err := newMemoryMeta()
assert.Nil(t, err)
schema := newTestSchema()
@ -463,7 +463,7 @@ func TestTryToSealSegment(t *testing.T) {
t.Run("normal seal with both segment & channel seal policy", func(t *testing.T) {
Params.Init()
mockAllocator := newMockAllocator()
meta, err := newMemoryMeta(mockAllocator)
meta, err := newMemoryMeta()
assert.Nil(t, err)
schema := newTestSchema()

33
internal/datacoord/server_test.go
View File

@ -44,7 +44,6 @@ import (
"github.com/milvus-io/milvus-proto/go-api/milvuspb"
"github.com/milvus-io/milvus-proto/go-api/schemapb"
"github.com/milvus-io/milvus/internal/common"
"github.com/milvus-io/milvus/internal/kv"
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
"github.com/milvus-io/milvus/internal/log"
"github.com/milvus-io/milvus/internal/mq/msgstream"
@ -2809,18 +2808,6 @@ func TestGetFlushState(t *testing.T) {
})
}
type mockTxnKVext struct {
kv.MockTxnKV
}
func (m *mockTxnKVext) LoadWithPrefix(prefix string) ([]string, []string, error) {
return []string{}, []string{}, nil
}
func (m *mockTxnKVext) MultiSave(kvs map[string]string) error {
return errors.New("(testing only) injected error")
}
func TestDataCoordServer_SetSegmentState(t *testing.T) {
t.Run("normal case", func(t *testing.T) {
svr := newTestServer(t, nil)
@ -2862,27 +2849,11 @@ func TestDataCoordServer_SetSegmentState(t *testing.T) {
assert.EqualValues(t, commonpb.SegmentState_Flushed, resp.States[0].State)
})
t.Run("dataCoord meta set state error", func(t *testing.T) {
meta, err := newMeta(context.TODO(), &mockTxnKVext{}, "")
t.Run("dataCoord meta set state not exists", func(t *testing.T) {
meta, err := newMemoryMeta()
assert.NoError(t, err)
svr := newTestServerWithMeta(t, nil, meta)
defer closeTestServer(t, svr)
segment := &datapb.SegmentInfo{
ID: 1000,
CollectionID: 100,
PartitionID: 0,
InsertChannel: "c1",
NumOfRows: 0,
State: commonpb.SegmentState_Growing,
StartPosition: &internalpb.MsgPosition{
ChannelName: "c1",
MsgID: []byte{},
MsgGroup: "",
Timestamp: 0,
},
}
err2 := svr.meta.AddSegment(NewSegmentInfo(segment))
assert.NotNil(t, err2)
// Set segment state.
svr.SetSegmentState(context.TODO(), &datapb.SetSegmentStateRequest{
SegmentId: 1000,

4
internal/kv/kv.go
View File

@ -50,7 +50,7 @@ type BaseKV interface {
Close()
}
//go:generate mockery --name=TxnKV
//go:generate mockery --name=TxnKV --with-expecter
// TxnKV contains extra txn operations of kv. The extra operations is transactional.
type TxnKV interface {
BaseKV
@ -78,7 +78,7 @@ type MetaKv interface {
CompareVersionAndSwap(key string, version int64, target string, opts ...clientv3.OpOption) (bool, error)
}
//go:generate mockery --name=SnapShotKV
//go:generate mockery --name=SnapShotKV --with-expecter
// SnapShotKV is TxnKV for snapshot data. It must save timestamp.
type SnapShotKV interface {
Save(key string, value string, ts typeutil.Timestamp) error

173
internal/kv/mock_kv.go
View File

@ -1,173 +0,0 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package kv
import (
"errors"
"strings"
"sync"
"github.com/milvus-io/milvus/internal/log"
clientv3 "go.etcd.io/etcd/client/v3"
"go.uber.org/zap"
)
type MockBaseKV struct {
InMemKv sync.Map
}
func (m *MockBaseKV) Load(key string) (string, error) {
log.Debug("doing load", zap.String("key", key))
if val, ok := m.InMemKv.Load(key); ok {
return val.(string), nil
}
return "", nil
}
func (m *MockBaseKV) MultiLoad(keys []string) ([]string, error) {
panic("not implemented") // TODO: Implement
}
func (m *MockBaseKV) LoadWithPrefix(key string) ([]string, []string, error) {
panic("not implemented") // TODO: Implement
}
func (m *MockBaseKV) Save(key string, value string) error {
m.InMemKv.Store(key, value)
log.Debug("doing Save", zap.String("key", key))
return nil
}
func (m *MockBaseKV) MultiSave(kvs map[string]string) error {
panic("not implemented") // TODO: Implement
}
func (m *MockBaseKV) Remove(key string) error {
m.InMemKv.Delete(key)
log.Debug("doing Remove", zap.String("key", key))
return nil
}
func (m *MockBaseKV) MultiRemove(keys []string) error {
panic("not implemented") // TODO: Implement
}
func (m *MockBaseKV) RemoveWithPrefix(key string) error {
panic("not implemented") // TODO: Implement
}
func (m *MockBaseKV) Close() {
panic("not implemented") // TODO: Implement
}
type MockTxnKV struct {
MockBaseKV
}
func (m *MockTxnKV) MultiSaveAndRemove(saves map[string]string, removals []string) error {
panic("not implemented") // TODO: Implement
}
func (m *MockTxnKV) MultiRemoveWithPrefix(keys []string) error {
panic("not implemented") // TODO: Implement
}
func (m *MockTxnKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string) error {
panic("not implemented") // TODO: Implement
}
type MockMetaKV struct {
MockTxnKV
LoadWithPrefixMockErr bool
SaveMockErr bool
}
func (m *MockMetaKV) GetPath(key string) string {
panic("not implemented") // TODO: Implement
}
func (m *MockMetaKV) LoadWithPrefix(prefix string) ([]string, []string, error) {
if m.LoadWithPrefixMockErr {
return nil, nil, errors.New("mock err")
}
keys := make([]string, 0)
values := make([]string, 0)
m.InMemKv.Range(func(key, value interface{}) bool {
if strings.HasPrefix(key.(string), prefix) {
keys = append(keys, key.(string))
values = append(values, value.(string))
}
return true
})
return keys, values, nil
}
func (m *MockMetaKV) LoadWithPrefix2(key string) ([]string, []string, []int64, error) {
panic("not implemented") // TODO: Implement
}
func (m *MockMetaKV) LoadWithRevisionAndVersions(key string) ([]string, []string, []int64, int64, error) {
panic("not implemented") // TODO: Implement
}
func (m *MockMetaKV) LoadWithRevision(key string) ([]string, []string, int64, error) {
panic("not implemented") // TODO: Implement
}
func (m *MockMetaKV) Watch(key string) clientv3.WatchChan {
panic("not implemented") // TODO: Implement
}
func (m *MockMetaKV) WatchWithPrefix(key string) clientv3.WatchChan {
panic("not implemented") // TODO: Implement
}
func (m *MockMetaKV) WatchWithRevision(key string, revision int64) clientv3.WatchChan {
panic("not implemented") // TODO: Implement
}
func (m *MockMetaKV) SaveWithLease(key, value string, id clientv3.LeaseID) error {
m.InMemKv.Store(key, value)
log.Debug("Doing SaveWithLease", zap.String("key", key))
return nil
}
func (m *MockMetaKV) SaveWithIgnoreLease(key, value string) error {
if m.SaveMockErr {
return errors.New("mock error")
}
m.InMemKv.Store(key, value)
log.Debug("Doing SaveWithIgnoreLease", zap.String("key", key))
return nil
}
func (m *MockMetaKV) Grant(ttl int64) (id clientv3.LeaseID, err error) {
return 1, nil
}
func (m *MockMetaKV) KeepAlive(id clientv3.LeaseID) (<-chan *clientv3.LeaseKeepAliveResponse, error) {
panic("not implemented") // TODO: Implement
}
func (m *MockMetaKV) CompareValueAndSwap(key, value, target string, opts ...clientv3.OpOption) (bool, error) {
panic("not implemented") // TODO: Implement
}
func (m *MockMetaKV) CompareVersionAndSwap(key string, version int64, target string, opts ...clientv3.OpOption) (bool, error) {
panic("not implemented") // TODO: Implement
}

129
internal/kv/mock_kv_test.go
View File

@ -1,129 +0,0 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package kv
import (
"sync"
"testing"
"github.com/stretchr/testify/assert"
clientv3 "go.etcd.io/etcd/client/v3"
)
const testKey = "key"
const testValue = "value"
func TestMockKV_MetaKV(t *testing.T) {
mockKv := &MockMetaKV{}
mockKv.InMemKv = sync.Map{}
var err error
value, err := mockKv.Load(testKey)
assert.Equal(t, "", value)
assert.NoError(t, err)
assert.Panics(t, func() {
mockKv.MultiLoad([]string{testKey})
})
_, _, err = mockKv.LoadWithPrefix(testKey)
assert.NoError(t, err)
assert.NoError(t, mockKv.Save(testKey, testValue))
assert.Panics(t, func() {
mockKv.MultiSave(map[string]string{testKey: testValue})
})
assert.NoError(t, mockKv.Remove(testKey))
assert.Panics(t, func() {
mockKv.MultiRemove([]string{testKey})
})
assert.Panics(t, func() {
mockKv.RemoveWithPrefix(testKey)
})
assert.Panics(t, func() {
mockKv.Close()
})
assert.Panics(t, func() {
mockKv.MultiSaveAndRemove(map[string]string{testKey: testValue}, []string{testKey})
})
assert.Panics(t, func() {
mockKv.MultiRemoveWithPrefix([]string{testKey})
})
assert.Panics(t, func() {
mockKv.MultiSaveAndRemoveWithPrefix(map[string]string{testKey: testValue}, []string{testKey})
})
assert.Panics(t, func() {
mockKv.GetPath(testKey)
})
assert.Panics(t, func() {
mockKv.LoadWithPrefix2(testKey)
})
assert.Panics(t, func() {
mockKv.LoadWithPrefix2(testKey)
})
assert.Panics(t, func() {
mockKv.LoadWithRevisionAndVersions(testKey)
})
assert.Panics(t, func() {
mockKv.LoadWithRevision(testKey)
})
assert.Panics(t, func() {
mockKv.Watch(testKey)
})
assert.Panics(t, func() {
mockKv.WatchWithPrefix(testKey)
})
assert.Panics(t, func() {
mockKv.WatchWithRevision(testKey, 100)
})
err = mockKv.SaveWithLease(testKey, testValue, 100)
assert.NoError(t, err)
err = mockKv.SaveWithIgnoreLease(testKey, testValue)
assert.NoError(t, err)
leaseID, err := mockKv.Grant(100)
assert.Equal(t, clientv3.LeaseID(1), leaseID)
assert.NoError(t, err)
assert.Panics(t, func() {
mockKv.KeepAlive(100)
})
assert.Panics(t, func() {
mockKv.CompareValueAndSwap(testKey, testValue, testValue)
})
assert.Panics(t, func() {
mockKv.CompareVersionAndSwap(testKey, 100, testKey)
})
}

59
internal/kv/mock_txn_kv.go
View File

@ -1,59 +0,0 @@
package kv
type TxnKVMock struct {
TxnKV
SaveF func(key, value string) error
RemoveF func(key string) error
}
func (m TxnKVMock) Load(key string) (string, error) {
panic("implement me")
}
func (m TxnKVMock) MultiLoad(keys []string) ([]string, error) {
panic("implement me")
}
func (m TxnKVMock) LoadWithPrefix(key string) ([]string, []string, error) {
panic("implement me")
}
func (m TxnKVMock) Save(key, value string) error {
return m.SaveF(key, value)
}
func (m TxnKVMock) MultiSave(kvs map[string]string) error {
panic("implement me")
}
func (m TxnKVMock) Remove(key string) error {
return m.RemoveF(key)
}
func (m TxnKVMock) MultiRemove(keys []string) error {
panic("implement me")
}
func (m TxnKVMock) RemoveWithPrefix(key string) error {
panic("implement me")
}
func (m TxnKVMock) Close() {
panic("implement me")
}
func (m TxnKVMock) MultiSaveAndRemove(saves map[string]string, removals []string) error {
panic("implement me")
}
func (m TxnKVMock) MultiRemoveWithPrefix(keys []string) error {
panic("implement me")
}
func (m TxnKVMock) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string) error {
panic("implement me")
}
func NewMockTxnKV() *TxnKVMock {
return &TxnKVMock{}
}

130
internal/kv/mocks/SnapShotKV.go
View File

@ -9,6 +9,14 @@ type SnapShotKV struct {
mock.Mock
}
type SnapShotKV_Expecter struct {
mock *mock.Mock
}
func (_m *SnapShotKV) EXPECT() *SnapShotKV_Expecter {
return &SnapShotKV_Expecter{mock: &_m.Mock}
}
// Load provides a mock function with given fields: key, ts
func (_m *SnapShotKV) Load(key string, ts uint64) (string, error) {
ret := _m.Called(key, ts)
@ -30,6 +38,30 @@ func (_m *SnapShotKV) Load(key string, ts uint64) (string, error) {
return r0, r1
}
// SnapShotKV_Load_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Load'
type SnapShotKV_Load_Call struct {
*mock.Call
}
// Load is a helper method to define mock.On call
// - key string
// - ts uint64
func (_e *SnapShotKV_Expecter) Load(key interface{}, ts interface{}) *SnapShotKV_Load_Call {
return &SnapShotKV_Load_Call{Call: _e.mock.On("Load", key, ts)}
}
func (_c *SnapShotKV_Load_Call) Run(run func(key string, ts uint64)) *SnapShotKV_Load_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(string), args[1].(uint64))
})
return _c
}
func (_c *SnapShotKV_Load_Call) Return(_a0 string, _a1 error) *SnapShotKV_Load_Call {
_c.Call.Return(_a0, _a1)
return _c
}
// LoadWithPrefix provides a mock function with given fields: key, ts
func (_m *SnapShotKV) LoadWithPrefix(key string, ts uint64) ([]string, []string, error) {
ret := _m.Called(key, ts)
@ -62,6 +94,30 @@ func (_m *SnapShotKV) LoadWithPrefix(key string, ts uint64) ([]string, []string,
return r0, r1, r2
}
// SnapShotKV_LoadWithPrefix_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'LoadWithPrefix'
type SnapShotKV_LoadWithPrefix_Call struct {
*mock.Call
}
// LoadWithPrefix is a helper method to define mock.On call
// - key string
// - ts uint64
func (_e *SnapShotKV_Expecter) LoadWithPrefix(key interface{}, ts interface{}) *SnapShotKV_LoadWithPrefix_Call {
return &SnapShotKV_LoadWithPrefix_Call{Call: _e.mock.On("LoadWithPrefix", key, ts)}
}
func (_c *SnapShotKV_LoadWithPrefix_Call) Run(run func(key string, ts uint64)) *SnapShotKV_LoadWithPrefix_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(string), args[1].(uint64))
})
return _c
}
func (_c *SnapShotKV_LoadWithPrefix_Call) Return(_a0 []string, _a1 []string, _a2 error) *SnapShotKV_LoadWithPrefix_Call {
_c.Call.Return(_a0, _a1, _a2)
return _c
}
// MultiSave provides a mock function with given fields: kvs, ts
func (_m *SnapShotKV) MultiSave(kvs map[string]string, ts uint64) error {
ret := _m.Called(kvs, ts)
@ -76,6 +132,30 @@ func (_m *SnapShotKV) MultiSave(kvs map[string]string, ts uint64) error {
return r0
}
// SnapShotKV_MultiSave_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'MultiSave'
type SnapShotKV_MultiSave_Call struct {
*mock.Call
}
// MultiSave is a helper method to define mock.On call
// - kvs map[string]string
// - ts uint64
func (_e *SnapShotKV_Expecter) MultiSave(kvs interface{}, ts interface{}) *SnapShotKV_MultiSave_Call {
return &SnapShotKV_MultiSave_Call{Call: _e.mock.On("MultiSave", kvs, ts)}
}
func (_c *SnapShotKV_MultiSave_Call) Run(run func(kvs map[string]string, ts uint64)) *SnapShotKV_MultiSave_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(map[string]string), args[1].(uint64))
})
return _c
}
func (_c *SnapShotKV_MultiSave_Call) Return(_a0 error) *SnapShotKV_MultiSave_Call {
_c.Call.Return(_a0)
return _c
}
// MultiSaveAndRemoveWithPrefix provides a mock function with given fields: saves, removals, ts
func (_m *SnapShotKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, removals []string, ts uint64) error {
ret := _m.Called(saves, removals, ts)
@ -90,6 +170,31 @@ func (_m *SnapShotKV) MultiSaveAndRemoveWithPrefix(saves map[string]string, remo
return r0
}
// SnapShotKV_MultiSaveAndRemoveWithPrefix_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'MultiSaveAndRemoveWithPrefix'
type SnapShotKV_MultiSaveAndRemoveWithPrefix_Call struct {
*mock.Call
}
// MultiSaveAndRemoveWithPrefix is a helper method to define mock.On call
// - saves map[string]string
// - removals []string
// - ts uint64
func (_e *SnapShotKV_Expecter) MultiSaveAndRemoveWithPrefix(saves interface{}, removals interface{}, ts interface{}) *SnapShotKV_MultiSaveAndRemoveWithPrefix_Call {
return &SnapShotKV_MultiSaveAndRemoveWithPrefix_Call{Call: _e.mock.On("MultiSaveAndRemoveWithPrefix", saves, removals, ts)}
}
func (_c *SnapShotKV_MultiSaveAndRemoveWithPrefix_Call) Run(run func(saves map[string]string, removals []string, ts uint64)) *SnapShotKV_MultiSaveAndRemoveWithPrefix_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(map[string]string), args[1].([]string), args[2].(uint64))
})
return _c
}
func (_c *SnapShotKV_MultiSaveAndRemoveWithPrefix_Call) Return(_a0 error) *SnapShotKV_MultiSaveAndRemoveWithPrefix_Call {
_c.Call.Return(_a0)
return _c
}
// Save provides a mock function with given fields: key, value, ts
func (_m *SnapShotKV) Save(key string, value string, ts uint64) error {
ret := _m.Called(key, value, ts)
@ -104,6 +209,31 @@ func (_m *SnapShotKV) Save(key string, value string, ts uint64) error {
return r0
}
// SnapShotKV_Save_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Save'
type SnapShotKV_Save_Call struct {
*mock.Call
}
// Save is a helper method to define mock.On call
// - key string
// - value string
// - ts uint64
func (_e *SnapShotKV_Expecter) Save(key interface{}, value interface{}, ts interface{}) *SnapShotKV_Save_Call {
return &SnapShotKV_Save_Call{Call: _e.mock.On("Save", key, value, ts)}
}
func (_c *SnapShotKV_Save_Call) Run(run func(key string, value string, ts uint64)) *SnapShotKV_Save_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(string), args[1].(string), args[2].(uint64))
})
return _c
}
func (_c *SnapShotKV_Save_Call) Return(_a0 error) *SnapShotKV_Save_Call {
_c.Call.Return(_a0)
return _c
}
type mockConstructorTestingTNewSnapShotKV interface {
mock.TestingT
Cleanup(func())

4
internal/rootcoord/import_manager.go
View File

@ -66,7 +66,7 @@ var flipTaskStateInterval = 15 * 1000
// importManager manager for import tasks
type importManager struct {
ctx context.Context // reserved
taskStore kv.MetaKv // Persistent task info storage.
taskStore kv.TxnKV // Persistent task info storage.
busyNodes map[int64]int64 // Set of all current working DataNode IDs and related task create timestamp.
// TODO: Make pendingTask a map to improve look up performance.
@ -89,7 +89,7 @@ type importManager struct {
}
// newImportManager helper function to create a importManager
func newImportManager(ctx context.Context, client kv.MetaKv,
func newImportManager(ctx context.Context, client kv.TxnKV,
idAlloc func(count uint32) (typeutil.UniqueID, typeutil.UniqueID, error),
importService func(ctx context.Context, req *datapb.ImportTaskRequest) (*datapb.ImportTaskResponse, error),
markSegmentsDropped func(ctx context.Context, segIDs []typeutil.UniqueID) (*commonpb.Status, error),

71
internal/rootcoord/import_manager_test.go
View File

@ -26,19 +26,17 @@ import (
"github.com/golang/protobuf/proto"
"github.com/milvus-io/milvus-proto/go-api/commonpb"
"github.com/milvus-io/milvus-proto/go-api/milvuspb"
"github.com/milvus-io/milvus/internal/kv"
memkv "github.com/milvus-io/milvus/internal/kv/mem"
"github.com/milvus-io/milvus/internal/kv/mocks"
"github.com/milvus-io/milvus/internal/proto/datapb"
"github.com/milvus-io/milvus/internal/proto/indexpb"
"github.com/milvus-io/milvus/internal/proto/rootcoordpb"
"github.com/milvus-io/milvus/internal/util/importutil"
"github.com/milvus-io/milvus/internal/util/typeutil"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
)
type customKV struct {
kv.MockMetaKV
}
func TestImportManager_NewImportManager(t *testing.T) {
var countLock sync.RWMutex
var globalCount = typeutil.UniqueID(0)
@ -54,8 +52,7 @@ func TestImportManager_NewImportManager(t *testing.T) {
Params.RootCoordCfg.ImportTaskRetention = 200
checkPendingTasksInterval = 100
cleanUpLoopInterval = 100
mockKv := &kv.MockMetaKV{}
mockKv.InMemKv = sync.Map{}
mockKv := memkv.NewMemoryKV()
ti1 := &datapb.ImportTaskInfo{
Id: 100,
State: &datapb.ImportTaskState{
@ -136,13 +133,13 @@ func TestImportManager_NewImportManager(t *testing.T) {
wg.Add(1)
t.Run("importManager init fail because of loadFromTaskStore fail", func(t *testing.T) {
defer wg.Done()
mockTxnKV := &mocks.TxnKV{}
mockTxnKV.EXPECT().LoadWithPrefix(mock.Anything).Return(nil, nil, errors.New("mock error"))
ctx, cancel := context.WithTimeout(context.Background(), 1*time.Nanosecond)
defer cancel()
mgr := newImportManager(ctx, mockKv, idAlloc, callImportServiceFn, callMarkSegmentsDropped, nil, nil, nil, nil)
mockKv.LoadWithPrefixMockErr = true
defer func() {
mockKv.LoadWithPrefixMockErr = false
}()
mgr := newImportManager(ctx, mockTxnKV, idAlloc, callImportServiceFn, callMarkSegmentsDropped, nil, nil, nil, nil)
assert.NotNil(t, mgr)
assert.Panics(t, func() {
mgr.init(context.TODO())
@ -152,13 +149,14 @@ func TestImportManager_NewImportManager(t *testing.T) {
wg.Add(1)
t.Run("sendOutTasks fail", func(t *testing.T) {
defer wg.Done()
mockTxnKV := &mocks.TxnKV{}
mockTxnKV.EXPECT().LoadWithPrefix(mock.Anything).Return(nil, nil, nil)
mockTxnKV.EXPECT().Save(mock.Anything, mock.Anything).Return(errors.New("mock save error"))
ctx, cancel := context.WithTimeout(context.Background(), 1*time.Nanosecond)
defer cancel()
mgr := newImportManager(ctx, mockKv, idAlloc, callImportServiceFn, callMarkSegmentsDropped, nil, nil, nil, nil)
mockKv.SaveMockErr = true
defer func() {
mockKv.SaveMockErr = false
}()
mgr := newImportManager(ctx, mockTxnKV, idAlloc, callImportServiceFn, callMarkSegmentsDropped, nil, nil, nil, nil)
assert.NotNil(t, mgr)
mgr.init(context.TODO())
})
@ -166,24 +164,23 @@ func TestImportManager_NewImportManager(t *testing.T) {
wg.Add(1)
t.Run("sendOutTasks fail", func(t *testing.T) {
defer wg.Done()
mockTxnKV := &mocks.TxnKV{}
mockTxnKV.EXPECT().LoadWithPrefix(mock.Anything).Return(nil, nil, nil)
ctx, cancel := context.WithTimeout(context.Background(), 1*time.Nanosecond)
defer cancel()
mgr := newImportManager(ctx, mockKv, idAlloc, callImportServiceFn, callMarkSegmentsDropped, nil, nil, nil, nil)
mgr := newImportManager(ctx, mockTxnKV, idAlloc, callImportServiceFn, callMarkSegmentsDropped, nil, nil, nil, nil)
assert.NotNil(t, mgr)
mgr.init(context.TODO())
func() {
mockKv.SaveMockErr = true
defer func() {
mockKv.SaveMockErr = false
}()
mockTxnKV.EXPECT().Save(mock.Anything, mock.Anything).Maybe().Return(errors.New("mock save error"))
mgr.sendOutTasks(context.TODO())
}()
func() {
mockTxnKV.EXPECT().Save(mock.Anything, mock.Anything).Maybe().Return(nil)
mockCallImportServiceErr = true
defer func() {
mockKv.SaveMockErr = false
}()
mgr.sendOutTasks(context.TODO())
}()
})
@ -254,8 +251,7 @@ func TestImportManager_TestSetImportTaskState(t *testing.T) {
Params.RootCoordCfg.ImportTaskRetention = 200
checkPendingTasksInterval = 100
cleanUpLoopInterval = 100
mockKv := &kv.MockMetaKV{}
mockKv.InMemKv = sync.Map{}
mockKv := memkv.NewMemoryKV()
ti1 := &datapb.ImportTaskInfo{
Id: 100,
State: &datapb.ImportTaskState{
@ -324,8 +320,7 @@ func TestImportManager_TestEtcdCleanUp(t *testing.T) {
Params.RootCoordCfg.ImportTaskRetention = 200
checkPendingTasksInterval = 100
cleanUpLoopInterval = 100
mockKv := &kv.MockMetaKV{}
mockKv.InMemKv = sync.Map{}
mockKv := memkv.NewMemoryKV()
ti1 := &datapb.ImportTaskInfo{
Id: 100,
State: &datapb.ImportTaskState{
@ -413,8 +408,7 @@ func TestImportManager_TestFlipTaskStateLoop(t *testing.T) {
Params.RootCoordCfg.ImportTaskRetention = 200
checkPendingTasksInterval = 100
cleanUpLoopInterval = 100
mockKv := &kv.MockMetaKV{}
mockKv.InMemKv = sync.Map{}
mockKv := memkv.NewMemoryKV()
ti1 := &datapb.ImportTaskInfo{
Id: 100,
State: &datapb.ImportTaskState{
@ -566,8 +560,7 @@ func TestImportManager_ImportJob(t *testing.T) {
}
Params.RootCoordCfg.ImportTaskSubPath = "test_import_task"
colID := int64(100)
mockKv := &kv.MockMetaKV{}
mockKv.InMemKv = sync.Map{}
mockKv := memkv.NewMemoryKV()
callMarkSegmentsDropped := func(ctx context.Context, segIDs []typeutil.UniqueID) (*commonpb.Status, error) {
return &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
@ -693,8 +686,7 @@ func TestImportManager_AllDataNodesBusy(t *testing.T) {
}
Params.RootCoordCfg.ImportTaskSubPath = "test_import_task"
colID := int64(100)
mockKv := &kv.MockMetaKV{}
mockKv.InMemKv = sync.Map{}
mockKv := memkv.NewMemoryKV()
rowReq := &milvuspb.ImportRequest{
CollectionName: "c1",
PartitionName: "p1",
@ -790,8 +782,7 @@ func TestImportManager_TaskState(t *testing.T) {
}
Params.RootCoordCfg.ImportTaskSubPath = "test_import_task"
colID := int64(100)
mockKv := &kv.MockMetaKV{}
mockKv.InMemKv = sync.Map{}
mockKv := memkv.NewMemoryKV()
importServiceFunc := func(ctx context.Context, req *datapb.ImportTaskRequest) (*datapb.ImportTaskResponse, error) {
return &datapb.ImportTaskResponse{
Status: &commonpb.Status{
@ -891,8 +882,7 @@ func TestImportManager_AllocFail(t *testing.T) {
}
Params.RootCoordCfg.ImportTaskSubPath = "test_import_task"
colID := int64(100)
mockKv := &kv.MockMetaKV{}
mockKv.InMemKv = sync.Map{}
mockKv := memkv.NewMemoryKV()
importServiceFunc := func(ctx context.Context, req *datapb.ImportTaskRequest) (*datapb.ImportTaskResponse, error) {
return &datapb.ImportTaskResponse{
Status: &commonpb.Status{
@ -931,8 +921,7 @@ func TestImportManager_ListAllTasks(t *testing.T) {
Params.RootCoordCfg.ImportTaskSubPath = "test_import_task"
colID := int64(100)
mockKv := &kv.MockMetaKV{}
mockKv.InMemKv = sync.Map{}
mockKv := memkv.NewMemoryKV()
// reject some tasks so there are 3 tasks left in pending list
fn := func(ctx context.Context, req *datapb.ImportTaskRequest) (*datapb.ImportTaskResponse, error) {

12
internal/rootcoord/mock_test.go
View File

@ -9,7 +9,6 @@ import (
"github.com/milvus-io/milvus-proto/go-api/commonpb"
"github.com/milvus-io/milvus-proto/go-api/milvuspb"
"github.com/milvus-io/milvus/internal/allocator"
"github.com/milvus-io/milvus/internal/kv"
"github.com/milvus-io/milvus/internal/log"
"github.com/milvus-io/milvus/internal/metastore/model"
"github.com/milvus-io/milvus/internal/mq/msgstream"
@ -250,17 +249,6 @@ func newMockTsoAllocator() *tso.MockAllocator {
return r
}
func newTxnKV() *kv.TxnKVMock {
r := kv.NewMockTxnKV()
r.SaveF = func(key, value string) error {
return nil
}
r.RemoveF = func(key string) error {
return nil
}
return r
}
type mockProxy struct {
types.Proxy
InvalidateCollectionMetaCacheFunc func(ctx context.Context, request *proxypb.InvalidateCollMetaCacheRequest) (*commonpb.Status, error)

11
internal/rootcoord/root_coord_test.go
View File

@ -13,7 +13,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/commonpb"
"github.com/milvus-io/milvus-proto/go-api/milvuspb"
"github.com/milvus-io/milvus/internal/allocator"
"github.com/milvus-io/milvus/internal/kv"
memkv "github.com/milvus-io/milvus/internal/kv/mem"
"github.com/milvus-io/milvus/internal/metastore/model"
"github.com/milvus-io/milvus/internal/proto/datapb"
"github.com/milvus-io/milvus/internal/proto/etcdpb"
@ -864,8 +864,7 @@ func TestCore_Import(t *testing.T) {
}
func TestCore_GetImportState(t *testing.T) {
mockKv := &kv.MockMetaKV{}
mockKv.InMemKv = sync.Map{}
mockKv := memkv.NewMemoryKV()
ti1 := &datapb.ImportTaskInfo{
Id: 100,
State: &datapb.ImportTaskState{
@ -914,8 +913,7 @@ func TestCore_GetImportState(t *testing.T) {
}
func TestCore_ListImportTasks(t *testing.T) {
mockKv := &kv.MockMetaKV{}
mockKv.InMemKv = sync.Map{}
mockKv := memkv.NewMemoryKV()
ti1 := &datapb.ImportTaskInfo{
Id: 100,
CollectionName: "collection-A",
@ -980,8 +978,7 @@ func TestCore_ReportImport(t *testing.T) {
globalCount++
return globalCount, 0, nil
}
mockKv := &kv.MockMetaKV{}
mockKv.InMemKv = sync.Map{}
mockKv := memkv.NewMemoryKV()
ti1 := &datapb.ImportTaskInfo{
Id: 100,
State: &datapb.ImportTaskState{