Remove redundant kv mocks (#25420)

Signed-off-by: sunby <bingyi.sun@zilliz.com>
Co-authored-by: sunby <bingyi.sun@zilliz.com>

internal/datacoord/cluster_test.go

@@ -24,11 +24,13 @@ import (
 	"github.com/cockroachdb/errors"
 	"github.com/golang/protobuf/proto"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 	"stathat.com/c/consistent"
 
 	"github.com/milvus-io/milvus/internal/kv"
 	etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
+	"github.com/milvus-io/milvus/internal/kv/mocks"
 	"github.com/milvus-io/milvus/internal/proto/datapb"
 	"github.com/milvus-io/milvus/internal/types"
 )
@@ -163,22 +165,13 @@ func TestClusterCreate(t *testing.T) {
 	t.Run("loadKv Fails", func(t *testing.T) {
 		defer kv.RemoveWithPrefix("")
 
-		fkv := &loadPrefixFailKV{WatchKV: kv}
+		metakv := mocks.NewWatchKV(t)
-		_, err := NewChannelManager(fkv, newMockHandler())
+		metakv.EXPECT().LoadWithPrefix(mock.Anything).Return(nil, nil, errors.New("failed"))
+		_, err := NewChannelManager(metakv, newMockHandler())
 		assert.Error(t, err)
 	})
 }
 
-// a mock kv that always fail when LoadWithPrefix
-type loadPrefixFailKV struct {
-	kv.WatchKV
-}
-
-// LoadWithPrefix override behavior
-func (kv *loadPrefixFailKV) LoadWithPrefix(key string) ([]string, []string, error) {
-	return []string{}, []string{}, errors.New("mocked fail")
-}
-
 func TestRegister(t *testing.T) {
 	kv := getWatchKV(t)
 	defer func() {

internal/datacoord/compaction_test.go

@@ -37,6 +37,8 @@ import (
 	"github.com/milvus-io/milvus/pkg/util/metautil"
 	"github.com/milvus-io/milvus/pkg/util/paramtable"
 	"github.com/milvus-io/milvus/pkg/util/typeutil"
+
+	mockkv "github.com/milvus-io/milvus/internal/kv/mocks"
 )
 
 func Test_compactionPlanHandler_execCompactionPlan(t *testing.T) {
@@ -341,8 +343,12 @@ func TestCompactionPlanHandler_handleMergeCompactionResult(t *testing.T) {
 
 	plans := map[int64]*compactionTask{1: task}
 
+	metakv := mockkv.NewMetaKv(t)
+	metakv.EXPECT().Save(mock.Anything, mock.Anything).Return(errors.New("failed")).Maybe()
+	metakv.EXPECT().MultiSave(mock.Anything).Return(errors.New("failed")).Maybe()
+	metakv.EXPECT().LoadWithPrefix(mock.Anything).Return(nil, nil, nil).Maybe()
 	errMeta := &meta{
-		catalog: &datacoord.Catalog{MetaKv: &saveFailKV{MetaKv: NewMetaMemoryKV()}},
+		catalog: &datacoord.Catalog{MetaKv: metakv},
 		segments: &SegmentsInfo{
 			map[int64]*SegmentInfo{
 				seg1.ID: {SegmentInfo: seg1},

internal/datacoord/index_meta_test.go

@@ -29,6 +29,7 @@ import (
 
 	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
 	"github.com/milvus-io/milvus/internal/kv/mocks"
+	mockkv "github.com/milvus-io/milvus/internal/kv/mocks"
 	"github.com/milvus-io/milvus/internal/metastore/kv/datacoord"
 	catalogmocks "github.com/milvus-io/milvus/internal/metastore/mocks"
 	"github.com/milvus-io/milvus/internal/metastore/model"
@@ -387,10 +388,14 @@ func TestMeta_GetIndexIDByName(t *testing.T) {
 			},
 		}
 	)
+	metakv := mockkv.NewMetaKv(t)
+	metakv.EXPECT().Save(mock.Anything, mock.Anything).Return(errors.New("failed")).Maybe()
+	metakv.EXPECT().MultiSave(mock.Anything).Return(errors.New("failed")).Maybe()
+	metakv.EXPECT().LoadWithPrefix(mock.Anything).Return(nil, nil, nil).Maybe()
 	m := &meta{
 		RWMutex:              sync.RWMutex{},
 		ctx:                  context.Background(),
-		catalog:              &datacoord.Catalog{MetaKv: &saveFailKV{}},
+		catalog:              &datacoord.Catalog{MetaKv: metakv},
 		indexes:              make(map[UniqueID]map[UniqueID]*model.Index),
 		buildID2SegmentIndex: make(map[UniqueID]*model.SegmentIndex),
 	}
@@ -445,10 +450,14 @@ func TestMeta_GetSegmentIndexState(t *testing.T) {
 			},
 		}
 	)
+	metakv := mockkv.NewMetaKv(t)
+	metakv.EXPECT().Save(mock.Anything, mock.Anything).Return(errors.New("failed")).Maybe()
+	metakv.EXPECT().MultiSave(mock.Anything).Return(errors.New("failed")).Maybe()
+	metakv.EXPECT().LoadWithPrefix(mock.Anything).Return(nil, nil, nil).Maybe()
 	m := &meta{
 		RWMutex:              sync.RWMutex{},
 		ctx:                  context.Background(),
-		catalog:              &datacoord.Catalog{MetaKv: &saveFailKV{}},
+		catalog:              &datacoord.Catalog{MetaKv: metakv},
 		indexes:              map[UniqueID]map[UniqueID]*model.Index{},
 		buildID2SegmentIndex: make(map[UniqueID]*model.SegmentIndex),
 		segments: &SegmentsInfo{
@@ -1122,8 +1131,11 @@ func TestMeta_FinishTask(t *testing.T) {
 	})
 
 	t.Run("fail", func(t *testing.T) {
+		metakv := mockkv.NewMetaKv(t)
+		metakv.EXPECT().Save(mock.Anything, mock.Anything).Return(errors.New("failed")).Maybe()
+		metakv.EXPECT().MultiSave(mock.Anything).Return(errors.New("failed")).Maybe()
 		m.catalog = &datacoord.Catalog{
-			MetaKv: &saveFailKV{},
+			MetaKv: metakv,
 		}
 		err := m.FinishTask(&indexpb.IndexTaskInfo{
 			BuildID:        buildID,

internal/datacoord/index_service_test.go

@@ -28,6 +28,7 @@ import (
 	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
 	"github.com/milvus-io/milvus-proto/go-api/v2/msgpb"
 	"github.com/milvus-io/milvus/internal/kv/mocks"
+	mockkv "github.com/milvus-io/milvus/internal/kv/mocks"
 	"github.com/milvus-io/milvus/internal/metastore/kv/datacoord"
 	catalogmocks "github.com/milvus-io/milvus/internal/metastore/mocks"
 	"github.com/milvus-io/milvus/internal/metastore/model"
@@ -135,8 +136,11 @@ func TestServer_CreateIndex(t *testing.T) {
 	})
 
 	t.Run("save index fail", func(t *testing.T) {
+		metakv := mockkv.NewMetaKv(t)
+		metakv.EXPECT().Save(mock.Anything, mock.Anything).Return(errors.New("failed")).Maybe()
+		metakv.EXPECT().MultiSave(mock.Anything).Return(errors.New("failed")).Maybe()
 		s.meta.indexes = map[UniqueID]map[UniqueID]*model.Index{}
-		s.meta.catalog = &datacoord.Catalog{MetaKv: &saveFailKV{}}
+		s.meta.catalog = &datacoord.Catalog{MetaKv: metakv}
 		req.IndexParams = []*commonpb.KeyValuePair{
 			{
 				Key:   common.IndexTypeKey,

internal/datacoord/meta_test.go

@@ -35,6 +35,8 @@ import (
 	"github.com/milvus-io/milvus/internal/mocks"
 	"github.com/milvus-io/milvus/internal/proto/datapb"
 	"github.com/milvus-io/milvus/pkg/common"
+
+	mockkv "github.com/milvus-io/milvus/internal/kv/mocks"
 )
 
 func TestMetaReloadFromKV(t *testing.T) {
@@ -269,17 +271,26 @@ func TestMeta_Basic(t *testing.T) {
 
 	t.Run("Test segment with kv fails", func(t *testing.T) {
 		// inject error for `Save`
-		memoryKV := NewMetaMemoryKV()
+		metakv := mockkv.NewMetaKv(t)
-		fkv := &saveFailKV{MetaKv: memoryKV}
+		metakv.EXPECT().Save(mock.Anything, mock.Anything).Return(errors.New("failed")).Maybe()
-		catalog := datacoord.NewCatalog(fkv, "", "")
+		metakv.EXPECT().MultiSave(mock.Anything).Return(errors.New("failed")).Maybe()
+		metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(nil).Maybe()
+		metakv.EXPECT().LoadWithPrefix(mock.Anything).Return(nil, nil, nil).Maybe()
+		catalog := datacoord.NewCatalog(metakv, "", "")
 		meta, err := newMeta(context.TODO(), catalog, nil)
 		assert.NoError(t, err)
 
 		err = meta.AddSegment(NewSegmentInfo(&datapb.SegmentInfo{}))
 		assert.Error(t, err)
 
-		fkv2 := &removeFailKV{MetaKv: memoryKV}
+		metakv2 := mockkv.NewMetaKv(t)
-		catalog = datacoord.NewCatalog(fkv2, "", "")
+		metakv2.EXPECT().Save(mock.Anything, mock.Anything).Return(nil).Maybe()
+		metakv2.EXPECT().MultiSave(mock.Anything).Return(nil).Maybe()
+		metakv2.EXPECT().Remove(mock.Anything).Return(errors.New("failed")).Maybe()
+		metakv2.EXPECT().MultiRemove(mock.Anything).Return(errors.New("failed")).Maybe()
+		metakv2.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(nil).Maybe()
+		metakv2.EXPECT().LoadWithPrefix(mock.Anything).Return(nil, nil, nil).Maybe()
+		catalog = datacoord.NewCatalog(metakv2, "", "")
 		meta, err = newMeta(context.TODO(), catalog, nil)
 		assert.NoError(t, err)
 		// nil, since no segment yet
@@ -292,7 +303,7 @@ func TestMeta_Basic(t *testing.T) {
 		err = meta.DropSegment(0)
 		assert.Error(t, err)
 
-		catalog = datacoord.NewCatalog(fkv, "", "")
+		catalog = datacoord.NewCatalog(metakv, "", "")
 		meta, err = newMeta(context.TODO(), catalog, nil)
 		assert.NoError(t, err)
 	})
@@ -484,9 +495,12 @@ func TestUpdateFlushSegmentsInfo(t *testing.T) {
 	})
 
 	t.Run("test save etcd failed", func(t *testing.T) {
-		kv := NewMetaMemoryKV()
+		metakv := mockkv.NewMetaKv(t)
-		failedKv := &saveFailKV{kv}
+		metakv.EXPECT().Save(mock.Anything, mock.Anything).Return(errors.New("mocked fail")).Maybe()
-		catalog := datacoord.NewCatalog(failedKv, "", "")
+		metakv.EXPECT().MultiSave(mock.Anything).Return(errors.New("mocked fail")).Maybe()
+		metakv.EXPECT().WalkWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(nil).Maybe()
+		metakv.EXPECT().LoadWithPrefix(mock.Anything).Return(nil, nil, nil).Maybe()
+		catalog := datacoord.NewCatalog(metakv, "", "")
 		meta, err := newMeta(context.TODO(), catalog, nil)
 		assert.NoError(t, err)
 

internal/datacoord/mock_test.go

@@ -29,7 +29,6 @@ import (
 	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
 	"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
 	"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
-	"github.com/milvus-io/milvus/internal/kv"
 	memkv "github.com/milvus-io/milvus/internal/kv/mem"
 	"github.com/milvus-io/milvus/internal/metastore/kv/datacoord"
 	"github.com/milvus-io/milvus/internal/proto/datapb"
@@ -137,30 +136,6 @@ func (a *FailsAllocator) allocID(_ context.Context) (UniqueID, error) {
 	return 0, errors.New("always fail")
 }
 
-// a mock kv that always fail when do `Save`
-type saveFailKV struct{ kv.MetaKv }
-
-// Save override behavior
-func (kv *saveFailKV) Save(key, value string) error {
-	return errors.New("mocked fail")
-}
-
-func (kv *saveFailKV) MultiSave(kvs map[string]string) error {
-	return errors.New("mocked fail")
-}
-
-// a mock kv that always fail when do `Remove`
-type removeFailKV struct{ kv.MetaKv }
-
-// Remove override behavior, inject error
-func (kv *removeFailKV) MultiRemove(key []string) error {
-	return errors.New("mocked fail")
-}
-
-func (kv *removeFailKV) Remove(key string) error {
-	return errors.New("mocked fail")
-}
-
 func newMockAllocator() *MockAllocator {
 	return &MockAllocator{}
 }

internal/datacoord/segment_manager_test.go

@@ -23,10 +23,13 @@ import (
 	"testing"
 	"time"
 
+	"github.com/cockroachdb/errors"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
 
 	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
 	"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
+	mockkv "github.com/milvus-io/milvus/internal/kv/mocks"
 	"github.com/milvus-io/milvus/internal/metastore/kv/datacoord"
 	"github.com/milvus-io/milvus/internal/proto/datapb"
 	"github.com/milvus-io/milvus/pkg/util/metautil"
@@ -608,7 +611,6 @@ func TestTryToSealSegment(t *testing.T) {
 		Params.Init()
 		mockAllocator := newMockAllocator()
 		memoryKV := NewMetaMemoryKV()
-		fkv := &saveFailKV{MetaKv: memoryKV}
 		catalog := datacoord.NewCatalog(memoryKV, "", "")
 		meta, err := newMeta(context.TODO(), catalog, nil)
 		assert.NoError(t, err)
@@ -622,7 +624,11 @@ func TestTryToSealSegment(t *testing.T) {
 		assert.NoError(t, err)
 		assert.EqualValues(t, 1, len(allocations))
 
-		segmentManager.meta.catalog = &datacoord.Catalog{MetaKv: fkv}
+		metakv := mockkv.NewMetaKv(t)
+		metakv.EXPECT().Save(mock.Anything, mock.Anything).Return(errors.New("failed")).Maybe()
+		metakv.EXPECT().MultiSave(mock.Anything).Return(errors.New("failed")).Maybe()
+		metakv.EXPECT().LoadWithPrefix(mock.Anything).Return(nil, nil, nil).Maybe()
+		segmentManager.meta.catalog = &datacoord.Catalog{MetaKv: metakv}
 
 		ts, err := segmentManager.allocator.allocTimestamp(context.Background())
 		assert.NoError(t, err)
@@ -634,7 +640,6 @@ func TestTryToSealSegment(t *testing.T) {
 		Params.Init()
 		mockAllocator := newMockAllocator()
 		memoryKV := NewMetaMemoryKV()
-		fkv := &saveFailKV{MetaKv: memoryKV}
 		catalog := datacoord.NewCatalog(memoryKV, "", "")
 		meta, err := newMeta(context.TODO(), catalog, nil)
 		assert.NoError(t, err)
@@ -648,7 +653,11 @@ func TestTryToSealSegment(t *testing.T) {
 		assert.NoError(t, err)
 		assert.EqualValues(t, 1, len(allocations))
 
-		segmentManager.meta.catalog = &datacoord.Catalog{MetaKv: fkv}
+		metakv := mockkv.NewMetaKv(t)
+		metakv.EXPECT().Save(mock.Anything, mock.Anything).Return(errors.New("failed")).Maybe()
+		metakv.EXPECT().MultiSave(mock.Anything).Return(errors.New("failed")).Maybe()
+		metakv.EXPECT().LoadWithPrefix(mock.Anything).Return(nil, nil, nil).Maybe()
+		segmentManager.meta.catalog = &datacoord.Catalog{MetaKv: metakv}
 
 		ts, err := segmentManager.allocator.allocTimestamp(context.Background())
 		assert.NoError(t, err)

internal/metastore/kv/datacoord/kv_catalog_test.go

@@ -199,6 +199,7 @@ func (mc *MockedTxnKV) LoadWithPrefix(key string) ([]string, []string, error) {
 }
 
 func (mc *MockedTxnKV) MultiRemove(keys []string) error {
+
 	return mc.multiRemove(keys)
 }