package rootcoord

import (
	"context"
	"fmt"
	"math/rand"
	"strings"
	"sync"
	"testing"

	"github.com/cockroachdb/errors"
	"github.com/samber/lo"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
	"go.uber.org/atomic"
	"go.uber.org/zap"
	"google.golang.org/protobuf/proto"

	"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/json"
	"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/metastore"
	"github.com/milvus-io/milvus/internal/metastore/model"
	"github.com/milvus-io/milvus/pkg/v2/common"
	"github.com/milvus-io/milvus/pkg/v2/log"
	pb "github.com/milvus-io/milvus/pkg/v2/proto/etcdpb"
	"github.com/milvus-io/milvus/pkg/v2/proto/internalpb"
	"github.com/milvus-io/milvus/pkg/v2/util"
	"github.com/milvus-io/milvus/pkg/v2/util/crypto"
	"github.com/milvus-io/milvus/pkg/v2/util/etcd"
	"github.com/milvus-io/milvus/pkg/v2/util/funcutil"
	"github.com/milvus-io/milvus/pkg/v2/util/merr"
	"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
	"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
)

var (
	indexName = "idx"
	IndexID   = 1
	index     = model.Index{
		CollectionID: 1,
		IndexName:    indexName,
		IndexID:      1,
		FieldID:      1,
		IndexParams:  []*commonpb.KeyValuePair{{Key: common.IndexTypeKey, Value: "STL_SORT"}},
		IsDeleted:    true,
	}
)

const (
	testDb = 1000
)

func TestCatalog_ListCollections(t *testing.T) {
	ctx := context.Background()

	coll1 := &pb.CollectionInfo{
		ID:                         1,
		PartitionIDs:               []int64{100},
		PartitionNames:             []string{"0"},
		PartitionCreatedTimestamps: []uint64{1},
		Schema: &schemapb.CollectionSchema{
			Name: "c1",
			Fields: []*schemapb.FieldSchema{
				{
					FieldID: 1,
					Name:    "f1",
				},
			},
		},
	}

	coll2 := &pb.CollectionInfo{
		ID:   2,
		DbId: testDb,
		Schema: &schemapb.CollectionSchema{
			Name: "c1",
			Fields: []*schemapb.FieldSchema{
				{},
			},
		},
		State: pb.CollectionState_CollectionCreated,
	}

	coll3 := &pb.CollectionInfo{
		ID: 3,
		Schema: &schemapb.CollectionSchema{
			Name: "c1",
			Fields: []*schemapb.FieldSchema{
				{},
			},
		},
		State: pb.CollectionState_CollectionDropping,
	}

	targetErr := errors.New("fail")

	t.Run("load collection with prefix fail", func(t *testing.T) {
		kv := mocks.NewSnapShotKV(t)
		ts := uint64(1)
		kv.On("LoadWithPrefix", mock.Anything, CollectionMetaPrefix, ts).
			Return(nil, nil, targetErr)

		kc := NewCatalog(nil, kv)
		ret, err := kc.ListCollections(ctx, util.NonDBID, ts)
		assert.ErrorIs(t, err, targetErr)
		assert.Nil(t, ret)
	})

	t.Run("list partition fail", func(t *testing.T) {
		kv := mocks.NewSnapShotKV(t)
		ts := uint64(1)

		bColl, err := proto.Marshal(coll2)
		assert.NoError(t, err)
		kv.On("LoadWithPrefix", mock.Anything, CollectionMetaPrefix, ts).
			Return([]string{"key"}, []string{string(bColl)}, nil)
		kv.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(nil, nil, targetErr)
		kc := NewCatalog(nil, kv)

		ret, err := kc.ListCollections(ctx, util.NonDBID, ts)
		assert.ErrorIs(t, err, targetErr)
		assert.Nil(t, ret)
	})

	t.Run("list fields fail", func(t *testing.T) {
		kv := mocks.NewSnapShotKV(t)
		ts := uint64(1)

		bColl, err := proto.Marshal(coll2)
		assert.NoError(t, err)
		kv.On("LoadWithPrefix", mock.Anything, CollectionMetaPrefix, ts).
			Return([]string{"key"}, []string{string(bColl)}, nil)

		partitionMeta := &pb.PartitionInfo{}
		pm, err := proto.Marshal(partitionMeta)
		assert.NoError(t, err)

		kv.On("LoadWithPrefix", mock.Anything, mock.MatchedBy(
			func(prefix string) bool {
				return strings.HasPrefix(prefix, PartitionMetaPrefix)
			}), ts).
			Return([]string{"key"}, []string{string(pm)}, nil)

		kv.On("LoadWithPrefix", mock.Anything, mock.MatchedBy(
			func(prefix string) bool {
				return strings.HasPrefix(prefix, FieldMetaPrefix)
			}), ts).
			Return(nil, nil, targetErr)
		kc := NewCatalog(nil, kv)

		ret, err := kc.ListCollections(ctx, util.NonDBID, ts)
		assert.ErrorIs(t, err, targetErr)
		assert.Nil(t, ret)
	})

	t.Run("list collection ok for 210 version", func(t *testing.T) {
		kv := mocks.NewSnapShotKV(t)
		ts := uint64(1)

		bColl, err := proto.Marshal(coll1)
		assert.NoError(t, err)
		kv.On("LoadWithPrefix", mock.Anything, CollectionMetaPrefix, ts).
			Return([]string{"key"}, []string{string(bColl)}, nil)
		kv.On("MultiSaveAndRemove", mock.Anything, mock.Anything, mock.Anything, ts).Return(nil)
		kc := NewCatalog(nil, kv)

		ret, err := kc.ListCollections(ctx, util.NonDBID, ts)
		assert.NoError(t, err)
		assert.Equal(t, 1, len(ret))
		assert.Equal(t, coll1.ID, ret[0].CollectionID)
		assert.Equal(t, util.DefaultDBID, ret[0].DBID)
	})

	t.Run("list collection with db", func(t *testing.T) {
		kv := mocks.NewSnapShotKV(t)
		ts := uint64(1)

		bColl, err := proto.Marshal(coll2)
		assert.NoError(t, err)
		kv.On("LoadWithPrefix", mock.Anything, BuildDatabasePrefixWithDBID(testDb), ts).
			Return([]string{"key"}, []string{string(bColl)}, nil)

		partitionMeta := &pb.PartitionInfo{}
		pm, err := proto.Marshal(partitionMeta)
		assert.NoError(t, err)

		kv.On("LoadWithPrefix", mock.Anything, mock.MatchedBy(
			func(prefix string) bool {
				return strings.HasPrefix(prefix, PartitionMetaPrefix)
			}), ts).
			Return([]string{"rootcoord/partitions/1/1"}, []string{string(pm)}, nil)

		fieldMeta := &schemapb.FieldSchema{}
		fm, err := proto.Marshal(fieldMeta)
		assert.NoError(t, err)

		kv.On("LoadWithPrefix", mock.Anything, mock.MatchedBy(
			func(prefix string) bool {
				return strings.HasPrefix(prefix, FieldMetaPrefix)
			}), ts).
			Return([]string{"rootcoord/fields/1/1"}, []string{string(fm)}, nil)

		functionMeta := &schemapb.FunctionSchema{}
		fcm, err := proto.Marshal(functionMeta)
		assert.NoError(t, err)
		kv.On("LoadWithPrefix", mock.Anything, mock.MatchedBy(
			func(prefix string) bool {
				return strings.HasPrefix(prefix, FunctionMetaPrefix)
			}), ts).
			Return([]string{"rootcoord/functions/1/1"}, []string{string(fcm)}, nil)

		kv.On("LoadWithPrefix", mock.Anything, mock.MatchedBy(
			func(prefix string) bool {
				return strings.HasPrefix(prefix, StructArrayFieldMetaPrefix)
			}), ts).
			Return([]string{}, []string{}, nil)

		kc := NewCatalog(nil, kv)
		ret, err := kc.ListCollections(ctx, testDb, ts)
		assert.NoError(t, err)
		assert.NotNil(t, ret)
		assert.Equal(t, 1, len(ret))
		assert.Equal(t, int64(testDb), ret[0].DBID)
	})

	t.Run("list collection ok for the newest version", func(t *testing.T) {
		kv := mocks.NewSnapShotKV(t)
		ts := uint64(1)

		bColl, err := proto.Marshal(coll2)
		assert.NoError(t, err)

		aColl, err := proto.Marshal(coll3)
		assert.NoError(t, err)

		kv.On("LoadWithPrefix", mock.Anything, CollectionMetaPrefix, ts).
			Return([]string{"key", "key2"}, []string{string(bColl), string(aColl)}, nil)

		partitionMeta := &pb.PartitionInfo{}
		pm, err := proto.Marshal(partitionMeta)
		assert.NoError(t, err)

		kv.On("LoadWithPrefix", mock.Anything, mock.MatchedBy(
			func(prefix string) bool {
				return strings.HasPrefix(prefix, PartitionMetaPrefix)
			}), ts).
			Return([]string{"rootcoord/partitions/1/1"}, []string{string(pm)}, nil)

		fieldMeta := &schemapb.FieldSchema{}
		fm, err := proto.Marshal(fieldMeta)
		assert.NoError(t, err)

		kv.On("LoadWithPrefix", mock.Anything, mock.MatchedBy(
			func(prefix string) bool {
				return strings.HasPrefix(prefix, FieldMetaPrefix)
			}), ts).
			Return([]string{"rootcoord/fields/1/1"}, []string{string(fm)}, nil)

		functionMeta := &schemapb.FunctionSchema{}
		fcm, err := proto.Marshal(functionMeta)
		assert.NoError(t, err)
		kv.On("LoadWithPrefix", mock.Anything, mock.MatchedBy(
			func(prefix string) bool {
				return strings.HasPrefix(prefix, FunctionMetaPrefix)
			}), ts).
			Return([]string{"rootcoord/functions/1/1"}, []string{string(fcm)}, nil)

		kv.On("LoadWithPrefix", mock.Anything, mock.MatchedBy(
			func(prefix string) bool {
				return strings.HasPrefix(prefix, StructArrayFieldMetaPrefix)
			}), ts).
			Return([]string{}, []string{}, nil)

		kv.On("MultiSaveAndRemove", mock.Anything, mock.Anything, mock.Anything, ts).Return(nil)
		kc := NewCatalog(nil, kv)

		ret, err := kc.ListCollections(ctx, util.NonDBID, ts)
		assert.NoError(t, err)
		assert.NotNil(t, ret)
		assert.Equal(t, 2, len(ret))
		assert.Equal(t, int64(2), ret[0].CollectionID)
		assert.Equal(t, int64(3), ret[1].CollectionID)
	})
}

func TestCatalog_loadCollection(t *testing.T) {
	t.Run("load failed", func(t *testing.T) {
		ctx := context.Background()
		kv := mocks.NewSnapShotKV(t)
		kv.EXPECT().Load(mock.Anything, mock.Anything, mock.Anything).Return("", errors.New("mock"))
		kc := NewCatalog(nil, kv).(*Catalog)
		_, err := kc.loadCollection(ctx, testDb, 1, 0)
		assert.Error(t, err)
	})

	t.Run("load, not collection info", func(t *testing.T) {
		ctx := context.Background()
		kv := mocks.NewSnapShotKV(t)
		kv.EXPECT().Load(mock.Anything, mock.Anything, mock.Anything).Return("not in pb format", nil)
		kc := NewCatalog(nil, kv).(*Catalog)
		_, err := kc.loadCollection(ctx, testDb, 1, 0)
		assert.Error(t, err)
	})

	t.Run("load, normal collection info", func(t *testing.T) {
		ctx := context.Background()
		coll := &pb.CollectionInfo{ID: 1, DbId: util.DefaultDBID}
		value, err := proto.Marshal(coll)
		assert.NoError(t, err)
		kv := mocks.NewSnapShotKV(t)
		kv.EXPECT().Load(mock.Anything, mock.Anything, mock.Anything).Return(string(value), nil)
		kc := NewCatalog(nil, kv).(*Catalog)
		got, err := kc.loadCollection(ctx, util.DefaultDBID, 1, 0)
		assert.NoError(t, err)
		assert.Equal(t, got.GetID(), coll.GetID())
	})

	t.Run("load, nonDBID collection info", func(t *testing.T) {
		ctx := context.Background()
		coll := &pb.CollectionInfo{
			ID:     1,
			DbId:   util.NonDBID,
			Schema: &schemapb.CollectionSchema{},
		}
		value, err := proto.Marshal(coll)
		assert.NoError(t, err)
		kv := mocks.NewSnapShotKV(t)
		kv.EXPECT().Load(mock.Anything, mock.Anything, mock.Anything).Return(string(value), nil)
		kv.EXPECT().MultiSaveAndRemove(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
		kc := NewCatalog(nil, kv).(*Catalog)
		got, err := kc.loadCollection(ctx, util.NonDBID, 1, 0)
		assert.NoError(t, err)
		assert.Equal(t, got.GetID(), coll.GetID())
	})
}

func Test_partitionVersionAfter210(t *testing.T) {
	t.Run("after 210", func(t *testing.T) {
		coll := &pb.CollectionInfo{}
		assert.True(t, partitionVersionAfter210(coll))
	})

	t.Run("before 210", func(t *testing.T) {
		coll := &pb.CollectionInfo{
			PartitionIDs:               []int64{1},
			PartitionNames:             []string{"partition"},
			PartitionCreatedTimestamps: []uint64{2},
		}
		assert.False(t, partitionVersionAfter210(coll))
	})
}

func Test_partitionExist(t *testing.T) {
	t.Run("in partitions ids", func(t *testing.T) {
		coll := &pb.CollectionInfo{
			PartitionIDs: []int64{1},
		}
		assert.True(t, partitionExistByID(coll, 1))
	})

	t.Run("not exist", func(t *testing.T) {
		coll := &pb.CollectionInfo{}
		assert.False(t, partitionExistByID(coll, 1))
	})
}

func Test_partitionExistByName(t *testing.T) {
	t.Run("in partitions ids", func(t *testing.T) {
		coll := &pb.CollectionInfo{
			PartitionNames: []string{"partition"},
		}
		assert.True(t, partitionExistByName(coll, "partition"))
	})

	t.Run("not exist", func(t *testing.T) {
		coll := &pb.CollectionInfo{}
		assert.False(t, partitionExistByName(coll, "partition"))
	})
}

func TestCatalog_GetCollectionByID(t *testing.T) {
	ctx := context.TODO()
	ss := mocks.NewSnapShotKV(t)
	c := NewCatalog(nil, ss)

	ss.EXPECT().Load(mock.Anything, mock.Anything, mock.Anything).Return("", errors.New("load error")).Twice()
	coll, err := c.GetCollectionByID(ctx, 0, 1, 1)
	assert.Error(t, err)
	assert.Nil(t, coll)

	ss.EXPECT().Load(mock.Anything, mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, key string, ts uint64) (string, error) {
		collByte, err := proto.Marshal(&pb.CollectionInfo{
			ID: 1,
			Schema: &schemapb.CollectionSchema{
				Fields: []*schemapb.FieldSchema{
					{},
				},
			},
			PartitionIDs:               []int64{1, 2, 3},
			PartitionNames:             []string{"1", "2", "3"},
			PartitionCreatedTimestamps: []uint64{1, 2, 3},
		})
		require.NoError(t, err)
		return string(collByte), nil
	}).Once()
	ss.EXPECT().MultiSaveAndRemove(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)

	coll, err = c.GetCollectionByID(ctx, 0, 10000, 1)
	assert.NoError(t, err)
	assert.NotNil(t, coll)
	assert.Equal(t, util.DefaultDBID, coll.DBID)
}

func TestCatalog_CreatePartitionV2(t *testing.T) {
	t.Run("collection not exist", func(t *testing.T) {
		ctx := context.Background()
		snapshot := kv.NewMockSnapshotKV()
		snapshot.LoadFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) {
			return "", errors.New("mock")
		}
		kc := NewCatalog(nil, snapshot).(*Catalog)
		err := kc.CreatePartition(ctx, 0, &model.Partition{}, 0)
		assert.Error(t, err)
	})

	t.Run("partition version after 210", func(t *testing.T) {
		ctx := context.Background()

		coll := &pb.CollectionInfo{
			DbId: util.DefaultDBID,
		}
		value, err := proto.Marshal(coll)
		assert.NoError(t, err)

		snapshot := kv.NewMockSnapshotKV()
		snapshot.LoadFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) {
			return string(value), nil
		}
		snapshot.SaveFunc = func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error {
			return errors.New("mock")
		}

		kc := NewCatalog(nil, snapshot).(*Catalog)

		err = kc.CreatePartition(ctx, 0, &model.Partition{}, 0)
		assert.Error(t, err)

		snapshot.SaveFunc = func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error {
			return nil
		}
		err = kc.CreatePartition(ctx, 0, &model.Partition{}, 0)
		assert.NoError(t, err)
	})

	t.Run("partition version before 210, id exist", func(t *testing.T) {
		ctx := context.Background()

		partID := typeutil.UniqueID(1)
		coll := &pb.CollectionInfo{DbId: util.DefaultDBID, PartitionIDs: []int64{partID}}
		value, err := proto.Marshal(coll)
		assert.NoError(t, err)

		snapshot := kv.NewMockSnapshotKV()
		snapshot.LoadFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) {
			return string(value), nil
		}

		kc := NewCatalog(nil, snapshot).(*Catalog)

		err = kc.CreatePartition(ctx, 0, &model.Partition{PartitionID: partID}, 0)
		assert.Error(t, err)
	})

	t.Run("partition version before 210, name exist", func(t *testing.T) {
		ctx := context.Background()

		partition := "partition"
		coll := &pb.CollectionInfo{DbId: util.DefaultDBID, PartitionNames: []string{partition}}
		value, err := proto.Marshal(coll)
		assert.NoError(t, err)

		snapshot := kv.NewMockSnapshotKV()
		snapshot.LoadFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) {
			return string(value), nil
		}

		kc := NewCatalog(nil, snapshot).(*Catalog)

		err = kc.CreatePartition(ctx, 0, &model.Partition{PartitionName: partition}, 0)
		assert.Error(t, err)
	})

	t.Run("partition version before 210, not exist", func(t *testing.T) {
		ctx := context.Background()

		coll := &pb.CollectionInfo{
			DbId:                       util.DefaultDBID,
			PartitionNames:             []string{"partition"},
			PartitionIDs:               []int64{111},
			PartitionCreatedTimestamps: []uint64{111111},
		}
		value, err := proto.Marshal(coll)
		assert.NoError(t, err)

		snapshot := kv.NewMockSnapshotKV()
		snapshot.LoadFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) {
			return string(value), nil
		}
		snapshot.SaveFunc = func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error {
			return errors.New("mock")
		}

		kc := NewCatalog(nil, snapshot).(*Catalog)

		err = kc.CreatePartition(ctx, 0, &model.Partition{}, 0)
		assert.Error(t, err)

		snapshot.SaveFunc = func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error {
			return nil
		}
		err = kc.CreatePartition(ctx, 0, &model.Partition{}, 0)
		assert.NoError(t, err)
	})
}

func TestCatalog_CreateAliasV2(t *testing.T) {
	ctx := context.Background()

	snapshot := kv.NewMockSnapshotKV()
	snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
		return errors.New("mock")
	}

	kc := NewCatalog(nil, snapshot).(*Catalog)

	err := kc.CreateAlias(ctx, &model.Alias{}, 0)
	assert.Error(t, err)

	snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
		return nil
	}
	err = kc.CreateAlias(ctx, &model.Alias{}, 0)
	assert.NoError(t, err)
}

func TestCatalog_listPartitionsAfter210(t *testing.T) {
	t.Run("load failed", func(t *testing.T) {
		ctx := context.Background()

		snapshot := kv.NewMockSnapshotKV()
		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
			return nil, nil, errors.New("mock")
		}

		kc := NewCatalog(nil, snapshot).(*Catalog)

		_, err := kc.listPartitionsAfter210(ctx, 1, 0)
		assert.Error(t, err)
	})

	t.Run("not in pb format", func(t *testing.T) {
		ctx := context.Background()

		snapshot := kv.NewMockSnapshotKV()
		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
			return []string{"key"}, []string{"not in pb format"}, nil
		}

		kc := NewCatalog(nil, snapshot).(*Catalog)

		_, err := kc.listPartitionsAfter210(ctx, 1, 0)
		assert.Error(t, err)
	})

	t.Run("normal case", func(t *testing.T) {
		ctx := context.Background()

		partition := &pb.PartitionInfo{PartitionID: 100}
		value, err := proto.Marshal(partition)
		assert.NoError(t, err)

		snapshot := kv.NewMockSnapshotKV()
		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
			return []string{"key"}, []string{string(value)}, nil
		}

		kc := NewCatalog(nil, snapshot).(*Catalog)

		got, err := kc.listPartitionsAfter210(ctx, 1, 0)
		assert.NoError(t, err)
		assert.Equal(t, 1, len(got))
		assert.Equal(t, int64(100), got[0].PartitionID)
	})
}

func Test_fieldVersionAfter210(t *testing.T) {
	coll := &pb.CollectionInfo{Schema: &schemapb.CollectionSchema{}}
	assert.True(t, fieldVersionAfter210(coll))

	coll.Schema.Fields = []*schemapb.FieldSchema{{FieldID: 101}}
	assert.False(t, fieldVersionAfter210(coll))
}

func TestCatalog_listFieldsAfter210(t *testing.T) {
	t.Run("load failed", func(t *testing.T) {
		ctx := context.Background()

		snapshot := kv.NewMockSnapshotKV()
		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
			return nil, nil, errors.New("mock")
		}

		kc := NewCatalog(nil, snapshot).(*Catalog)

		_, err := kc.listFieldsAfter210(ctx, 1, 0)
		assert.Error(t, err)
	})

	t.Run("not in pb format", func(t *testing.T) {
		ctx := context.Background()

		snapshot := kv.NewMockSnapshotKV()
		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
			return []string{"key"}, []string{"not in pb format"}, nil
		}

		kc := NewCatalog(nil, snapshot).(*Catalog)

		_, err := kc.listFieldsAfter210(ctx, 1, 0)
		assert.Error(t, err)
	})

	t.Run("normal case", func(t *testing.T) {
		ctx := context.Background()

		field := &schemapb.FieldSchema{FieldID: 101}
		value, err := proto.Marshal(field)
		assert.NoError(t, err)

		snapshot := kv.NewMockSnapshotKV()
		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
			return []string{"key"}, []string{string(value)}, nil
		}

		kc := NewCatalog(nil, snapshot).(*Catalog)

		got, err := kc.listFieldsAfter210(ctx, 1, 0)
		assert.NoError(t, err)
		assert.Equal(t, 1, len(got))
		assert.Equal(t, int64(101), got[0].FieldID)
	})
}

func TestCatalog_AlterAliasV2(t *testing.T) {
	ctx := context.Background()

	snapshot := kv.NewMockSnapshotKV()
	snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
		return errors.New("mock")
	}

	kc := NewCatalog(nil, snapshot).(*Catalog)

	err := kc.AlterAlias(ctx, &model.Alias{}, 0)
	assert.Error(t, err)

	snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
		return nil
	}
	err = kc.AlterAlias(ctx, &model.Alias{}, 0)
	assert.NoError(t, err)
}

func Test_dropPartition(t *testing.T) {
	t.Run("nil, won't panic", func(t *testing.T) {
		dropPartition(nil, 1)
	})

	t.Run("not exist", func(t *testing.T) {
		coll := &pb.CollectionInfo{}
		dropPartition(coll, 100)
	})

	t.Run("in partition ids", func(t *testing.T) {
		coll := &pb.CollectionInfo{
			PartitionIDs:               []int64{100, 101, 102, 103, 104},
			PartitionNames:             []string{"0", "1", "2", "3", "4"},
			PartitionCreatedTimestamps: []uint64{1, 2, 3, 4, 5},
		}
		dropPartition(coll, 100)
		assert.Equal(t, 4, len(coll.GetPartitionIDs()))
		dropPartition(coll, 104)
		assert.Equal(t, 3, len(coll.GetPartitionIDs()))
		dropPartition(coll, 102)
		assert.Equal(t, 2, len(coll.GetPartitionIDs()))
		dropPartition(coll, 103)
		assert.Equal(t, 1, len(coll.GetPartitionIDs()))
		dropPartition(coll, 101)
		assert.Equal(t, 0, len(coll.GetPartitionIDs()))
	})
}

func TestCatalog_DropPartitionV2(t *testing.T) {
	t.Run("failed to load collection", func(t *testing.T) {
		ctx := context.Background()

		snapshot := mocks.NewSnapShotKV(t)
		snapshot.On("Load",
			mock.Anything, mock.Anything, mock.Anything).Return("not in codec format", nil)

		kc := NewCatalog(nil, snapshot).(*Catalog)

		err := kc.DropPartition(ctx, 0, 100, 101, 0)
		assert.Error(t, err)
	})

	t.Run("failed to load collection, no key found", func(t *testing.T) {
		ctx := context.Background()

		snapshot := mocks.NewSnapShotKV(t)
		snapshot.On("Load",
			mock.Anything, mock.Anything, mock.Anything).Return("", merr.WrapErrIoKeyNotFound("partition"))

		kc := NewCatalog(nil, snapshot).(*Catalog)

		err := kc.DropPartition(ctx, 0, 100, 101, 0)
		assert.NoError(t, err)
	})

	t.Run("partition version after 210", func(t *testing.T) {
		ctx := context.Background()

		coll := &pb.CollectionInfo{DbId: util.DefaultDBID}
		value, err := proto.Marshal(coll)
		assert.NoError(t, err)

		snapshot := kv.NewMockSnapshotKV()
		snapshot.LoadFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) {
			return string(value), nil
		}
		snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
			return errors.New("mock")
		}

		kc := NewCatalog(nil, snapshot).(*Catalog)

		err = kc.DropPartition(ctx, 0, 100, 101, 0)
		assert.Error(t, err)

		snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
			return nil
		}
		err = kc.DropPartition(ctx, 0, 100, 101, 0)
		assert.NoError(t, err)
	})

	t.Run("partition before 210", func(t *testing.T) {
		ctx := context.Background()

		coll := &pb.CollectionInfo{
			DbId:                       util.DefaultDBID,
			PartitionIDs:               []int64{101, 102},
			PartitionNames:             []string{"partition1", "partition2"},
			PartitionCreatedTimestamps: []uint64{101, 102},
		}
		value, err := proto.Marshal(coll)
		assert.NoError(t, err)

		snapshot := kv.NewMockSnapshotKV()
		snapshot.LoadFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) (string, error) {
			return string(value), nil
		}
		snapshot.SaveFunc = func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error {
			return errors.New("mock")
		}

		kc := NewCatalog(nil, snapshot).(*Catalog)

		err = kc.DropPartition(ctx, 0, 100, 101, 0)
		assert.Error(t, err)

		snapshot.SaveFunc = func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error {
			return nil
		}
		err = kc.DropPartition(ctx, 0, 100, 102, 0)
		assert.NoError(t, err)
	})
}

func TestCatalog_DropAliasV2(t *testing.T) {
	ctx := context.Background()

	snapshot := kv.NewMockSnapshotKV()
	snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
		return errors.New("mock")
	}

	kc := NewCatalog(nil, snapshot).(*Catalog)

	err := kc.DropAlias(ctx, testDb, "alias", 0)
	assert.Error(t, err)

	snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
		return nil
	}
	err = kc.DropAlias(ctx, testDb, "alias", 0)
	assert.NoError(t, err)
}

func TestCatalog_listAliasesBefore210(t *testing.T) {
	t.Run("load failed", func(t *testing.T) {
		ctx := context.Background()

		snapshot := kv.NewMockSnapshotKV()
		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
			return nil, nil, errors.New("mock")
		}

		kc := NewCatalog(nil, snapshot).(*Catalog)

		_, err := kc.listAliasesBefore210(ctx, 0)
		assert.Error(t, err)
	})

	t.Run("not in pb format", func(t *testing.T) {
		ctx := context.Background()

		snapshot := kv.NewMockSnapshotKV()
		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
			return []string{"key"}, []string{"not in pb format"}, nil
		}

		kc := NewCatalog(nil, snapshot).(*Catalog)

		_, err := kc.listAliasesBefore210(ctx, 0)
		assert.Error(t, err)
	})

	t.Run("normal case", func(t *testing.T) {
		ctx := context.Background()

		coll := &pb.CollectionInfo{ID: 100}
		value, err := proto.Marshal(coll)
		assert.NoError(t, err)

		snapshot := kv.NewMockSnapshotKV()
		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
			return []string{"key"}, []string{string(value)}, nil
		}

		kc := NewCatalog(nil, snapshot).(*Catalog)

		got, err := kc.listAliasesBefore210(ctx, 0)
		assert.NoError(t, err)
		assert.Equal(t, 1, len(got))
		assert.Equal(t, int64(100), got[0].CollectionID)
	})
}

func TestCatalog_listAliasesAfter210(t *testing.T) {
	t.Run("load failed", func(t *testing.T) {
		ctx := context.Background()

		snapshot := kv.NewMockSnapshotKV()
		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
			return nil, nil, errors.New("mock")
		}

		kc := NewCatalog(nil, snapshot).(*Catalog)

		_, err := kc.listAliasesAfter210WithDb(ctx, testDb, 0)
		assert.Error(t, err)
	})

	t.Run("not in pb format", func(t *testing.T) {
		ctx := context.Background()

		snapshot := kv.NewMockSnapshotKV()
		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
			return []string{"key"}, []string{"not in pb format"}, nil
		}

		kc := NewCatalog(nil, snapshot).(*Catalog)

		_, err := kc.listAliasesAfter210WithDb(ctx, testDb, 0)
		assert.Error(t, err)
	})

	t.Run("normal case", func(t *testing.T) {
		ctx := context.Background()

		coll := &pb.AliasInfo{CollectionId: 100}
		value, err := proto.Marshal(coll)
		assert.NoError(t, err)

		snapshot := kv.NewMockSnapshotKV()
		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
			return []string{"key"}, []string{string(value)}, nil
		}

		kc := NewCatalog(nil, snapshot).(*Catalog)

		got, err := kc.listAliasesAfter210WithDb(ctx, testDb, 0)
		assert.NoError(t, err)
		assert.Equal(t, 1, len(got))
		assert.Equal(t, int64(100), got[0].CollectionID)
	})
}

func TestCatalog_ListAliasesV2(t *testing.T) {
	t.Run("failed to list aliases before 210", func(t *testing.T) {
		ctx := context.Background()

		snapshot := kv.NewMockSnapshotKV()
		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
			return []string{"key"}, []string{"not in pb format"}, nil
		}

		kc := NewCatalog(nil, snapshot).(*Catalog)

		_, err := kc.ListAliases(ctx, testDb, 0)
		assert.Error(t, err)
	})

	t.Run("failed to list aliases after 210", func(t *testing.T) {
		ctx := context.Background()

		coll := &pb.CollectionInfo{ID: 100, ShardsNum: 50}
		value, err := proto.Marshal(coll)
		assert.NoError(t, err)

		snapshot := kv.NewMockSnapshotKV()
		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
			if key == AliasMetaPrefix {
				return nil, nil, errors.New("mock")
			}

			if strings.Contains(key, DatabaseMetaPrefix) {
				return nil, nil, errors.New("mock")
			}
			return []string{"key"}, []string{string(value)}, nil
		}

		kc := NewCatalog(nil, snapshot).(*Catalog)

		_, err = kc.ListAliases(ctx, util.NonDBID, 0)
		assert.Error(t, err)

		_, err = kc.ListAliases(ctx, testDb, 0)
		assert.Error(t, err)
	})

	t.Run("normal case", func(t *testing.T) {
		ctx := context.Background()

		alias := &pb.AliasInfo{CollectionId: 101, AliasName: "alias2"}
		value2, err := proto.Marshal(alias)
		assert.NoError(t, err)

		snapshot := kv.NewMockSnapshotKV()
		snapshot.LoadWithPrefixFunc = func(ctx context.Context, key string, ts typeutil.Timestamp) ([]string, []string, error) {
			dbStr := fmt.Sprintf("%d", testDb)
			if strings.Contains(key, dbStr) && strings.Contains(key, Aliases) {
				return []string{"key1"}, []string{string(value2)}, nil
			}
			return []string{}, []string{}, nil
		}

		kc := NewCatalog(nil, snapshot).(*Catalog)

		got, err := kc.ListAliases(ctx, testDb, 0)
		assert.NoError(t, err)
		assert.Equal(t, 1, len(got))
		assert.Equal(t, "alias2", got[0].Name)
	})
}

func Test_batchMultiSaveAndRemove(t *testing.T) {
	t.Run("failed to save", func(t *testing.T) {
		snapshot := kv.NewMockSnapshotKV()
		snapshot.MultiSaveFunc = func(ctx context.Context, kvs map[string]string, ts typeutil.Timestamp) error {
			return errors.New("error mock MultiSave")
		}
		saves := map[string]string{"k": "v"}
		err := batchMultiSaveAndRemove(context.TODO(), snapshot, paramtable.Get().MetaStoreCfg.MaxEtcdTxnNum.GetAsInt(), saves, []string{}, 0)
		assert.Error(t, err)
	})
	t.Run("failed to remove", func(t *testing.T) {
		snapshot := kv.NewMockSnapshotKV()
		snapshot.MultiSaveFunc = func(ctx context.Context, kvs map[string]string, ts typeutil.Timestamp) error {
			return nil
		}
		snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
			return errors.New("error mock MultiSaveAndRemove")
		}
		saves := map[string]string{"k": "v"}
		removals := []string{"prefix1", "prefix2"}
		err := batchMultiSaveAndRemove(context.TODO(), snapshot, paramtable.Get().MetaStoreCfg.MaxEtcdTxnNum.GetAsInt(), saves, removals, 0)
		assert.Error(t, err)
	})
	t.Run("normal case", func(t *testing.T) {
		snapshot := kv.NewMockSnapshotKV()
		snapshot.MultiSaveFunc = func(ctx context.Context, kvs map[string]string, ts typeutil.Timestamp) error {
			log.Info("multi save", zap.Any("len", len(kvs)), zap.Any("saves", kvs))
			return nil
		}
		snapshot.MultiSaveAndRemoveFunc = func(ctx context.Context, saves map[string]string, removals []string, ts typeutil.Timestamp) error {
			log.Info("multi save and remove with prefix", zap.Any("len of saves", len(saves)), zap.Any("len of removals", len(removals)),
				zap.Any("saves", saves), zap.Any("removals", removals))
			return nil
		}
		n := 400
		saves := map[string]string{}
		removals := make([]string, 0, n)
		for i := 0; i < n; i++ {
			saves[fmt.Sprintf("k%d", i)] = fmt.Sprintf("v%d", i)
			removals = append(removals, fmt.Sprintf("k%d", i))
		}
		err := batchMultiSaveAndRemove(context.TODO(), snapshot, paramtable.Get().MetaStoreCfg.MaxEtcdTxnNum.GetAsInt(), saves, removals, 0)
		assert.NoError(t, err)
	})
}

func TestCatalog_AlterCollection(t *testing.T) {
	t.Run("add", func(t *testing.T) {
		kc := NewCatalog(nil, nil)
		ctx := context.Background()
		err := kc.AlterCollection(ctx, nil, nil, metastore.ADD, 0, false)
		assert.Error(t, err)
	})

	t.Run("delete", func(t *testing.T) {
		kc := NewCatalog(nil, nil)
		ctx := context.Background()
		err := kc.AlterCollection(ctx, nil, nil, metastore.DELETE, 0, false)
		assert.Error(t, err)
	})

	t.Run("modify", func(t *testing.T) {
		snapshot := kv.NewMockSnapshotKV()
		kvs := map[string]string{}
		snapshot.SaveFunc = func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error {
			kvs[key] = value
			return nil
		}
		snapshot.MultiSaveFunc = func(ctx context.Context, saveKvs map[string]string, _ typeutil.Timestamp) error {
			for k, v := range saveKvs {
				kvs[k] = v
			}
			return nil
		}
		kc := NewCatalog(nil, snapshot).(*Catalog)
		ctx := context.Background()
		var collectionID int64 = 1
		oldC := &model.Collection{CollectionID: collectionID, State: pb.CollectionState_CollectionCreated}
		newC := &model.Collection{CollectionID: collectionID, State: pb.CollectionState_CollectionCreated, UpdateTimestamp: rand.Uint64()}
		err := kc.AlterCollection(ctx, oldC, newC, metastore.MODIFY, 0, true)
		assert.NoError(t, err)
		key := BuildCollectionKey(0, collectionID)
		value, ok := kvs[key]
		assert.True(t, ok)
		var collPb pb.CollectionInfo
		err = proto.Unmarshal([]byte(value), &collPb)
		assert.NoError(t, err)
		got := model.UnmarshalCollectionModel(&collPb)
		assert.Equal(t, pb.CollectionState_CollectionCreated, got.State)
		assert.Equal(t, newC.UpdateTimestamp, got.UpdateTimestamp)
	})

	t.Run("modify EnableDynamicField and SchemaVersion", func(t *testing.T) {
		snapshot := kv.NewMockSnapshotKV()
		kvs := map[string]string{}
		snapshot.SaveFunc = func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error {
			kvs[key] = value
			return nil
		}
		snapshot.MultiSaveFunc = func(ctx context.Context, saveKvs map[string]string, _ typeutil.Timestamp) error {
			for k, v := range saveKvs {
				kvs[k] = v
			}
			return nil
		}
		kc := NewCatalog(nil, snapshot).(*Catalog)
		ctx := context.Background()
		var collectionID int64 = 1
		oldC := &model.Collection{
			CollectionID:       collectionID,
			State:              pb.CollectionState_CollectionCreated,
			EnableDynamicField: false,
			SchemaVersion:      1,
		}
		newC := &model.Collection{
			CollectionID:       collectionID,
			State:              pb.CollectionState_CollectionCreated,
			EnableDynamicField: true,
			SchemaVersion:      2,
			UpdateTimestamp:    rand.Uint64(),
		}
		err := kc.AlterCollection(ctx, oldC, newC, metastore.MODIFY, 0, true)
		assert.NoError(t, err)
		key := BuildCollectionKey(0, collectionID)
		value, ok := kvs[key]
		assert.True(t, ok)
		var collPb pb.CollectionInfo
		err = proto.Unmarshal([]byte(value), &collPb)
		assert.NoError(t, err)
		got := model.UnmarshalCollectionModel(&collPb)
		assert.Equal(t, pb.CollectionState_CollectionCreated, got.State)
		assert.Equal(t, newC.UpdateTimestamp, got.UpdateTimestamp)
		assert.Equal(t, newC.EnableDynamicField, got.EnableDynamicField)
		assert.Equal(t, newC.SchemaVersion, got.SchemaVersion)
	})

	t.Run("modify, tenant id changed", func(t *testing.T) {
		kc := NewCatalog(nil, nil)
		ctx := context.Background()
		var collectionID int64 = 1
		oldC := &model.Collection{TenantID: "1", CollectionID: collectionID, State: pb.CollectionState_CollectionCreated}
		newC := &model.Collection{TenantID: "2", CollectionID: collectionID, State: pb.CollectionState_CollectionCreated}
		err := kc.AlterCollection(ctx, oldC, newC, metastore.MODIFY, 0, true)
		assert.Error(t, err)
	})

	t.Run("modify db name", func(t *testing.T) {
		var collectionID int64 = 1
		snapshot := kv.NewMockSnapshotKV()

		kc := NewCatalog(nil, snapshot).(*Catalog)
		ctx := context.Background()
		oldC := &model.Collection{DBID: 0, CollectionID: collectionID, State: pb.CollectionState_CollectionCreated}
		newC := &model.Collection{DBID: 1, CollectionID: collectionID, State: pb.CollectionState_CollectionCreated}
		err := kc.AlterCollection(ctx, oldC, newC, metastore.MODIFY, 0, true)
		assert.Error(t, err)
	})

	t.Run("modify 64 fields", func(t *testing.T) {
		var collectionID int64 = 1
		snapshot := kv.NewMockSnapshotKV()
		snapshot.MultiSaveFunc = func(ctx context.Context, saves map[string]string, ts typeutil.Timestamp) error {
			assert.LessOrEqual(t, len(saves), 64)
			return nil
		}

		kc := NewCatalog(nil, snapshot).(*Catalog)
		ctx := context.Background()
		// 2 system fields + 64 user fields
		fields := make([]*model.Field, 66)
		for i := range fields {
			fields[i] = &model.Field{
				FieldID: int64(i),
			}
		}
		oldC := &model.Collection{DBID: 0, CollectionID: collectionID, State: pb.CollectionState_CollectionCreated, Fields: fields}
		newC := &model.Collection{DBID: 0, CollectionID: collectionID, State: pb.CollectionState_CollectionCreated, Fields: fields}
		err := kc.AlterCollection(ctx, oldC, newC, metastore.MODIFY, 0, true)
		assert.NoError(t, err)
	})

	t.Run("modify function", func(t *testing.T) {
		var collectionID int64 = 1
		snapshot := kv.NewMockSnapshotKV()
		snapshot.MultiSaveFunc = func(ctx context.Context, saves map[string]string, ts typeutil.Timestamp) error {
			assert.LessOrEqual(t, len(saves), 64)
			return nil
		}

		kc := NewCatalog(nil, snapshot).(*Catalog)
		ctx := context.Background()

		functions := []*model.Function{
			{
				Name: "test_function",
			},
		}

		oldC := &model.Collection{DBID: 0, CollectionID: collectionID, State: pb.CollectionState_CollectionCreated}
		newC := &model.Collection{DBID: 0, CollectionID: collectionID, State: pb.CollectionState_CollectionCreated, Functions: functions}
		err := kc.AlterCollection(ctx, oldC, newC, metastore.MODIFY, 0, true)
		assert.NoError(t, err)
	})
}

func TestCatalog_AlterCollectionDB(t *testing.T) {
	snapshot := kv.NewMockSnapshotKV()
	kvs := map[string]string{}
	snapshot.MultiSaveFunc = func(ctx context.Context, saveKvs map[string]string, _ typeutil.Timestamp) error {
		for k, v := range saveKvs {
			kvs[k] = v
		}
		return nil
	}
	kc := NewCatalog(nil, snapshot).(*Catalog)
	ctx := context.Background()
	var collectionID int64 = 1
	t.Run("rename tencentid", func(t *testing.T) {
		oldC := &model.Collection{TenantID: "0", CollectionID: collectionID, State: pb.CollectionState_CollectionCreated, DBID: 0}
		newC := &model.Collection{TenantID: "1", CollectionID: collectionID, State: pb.CollectionState_CollectionCreated, DBID: 1}
		err := kc.AlterCollectionDB(ctx, oldC, newC, 0)
		assert.Error(t, err)
	})

	t.Run("modify db", func(t *testing.T) {
		oldC := &model.Collection{CollectionID: collectionID, State: pb.CollectionState_CollectionCreated, DBID: 0}
		newC := &model.Collection{CollectionID: collectionID, State: pb.CollectionState_CollectionCreated, DBID: 1}
		err := kc.AlterCollectionDB(ctx, oldC, newC, 0)
		assert.NoError(t, err)
	})
}

func TestCatalog_AlterPartition(t *testing.T) {
	t.Run("add", func(t *testing.T) {
		kc := NewCatalog(nil, nil)
		ctx := context.Background()
		err := kc.AlterPartition(ctx, testDb, nil, nil, metastore.ADD, 0)
		assert.Error(t, err)
	})

	t.Run("delete", func(t *testing.T) {
		kc := NewCatalog(nil, nil)
		ctx := context.Background()
		err := kc.AlterPartition(ctx, testDb, nil, nil, metastore.DELETE, 0)
		assert.Error(t, err)
	})

	t.Run("modify", func(t *testing.T) {
		snapshot := kv.NewMockSnapshotKV()
		kvs := map[string]string{}
		snapshot.SaveFunc = func(ctx context.Context, key string, value string, ts typeutil.Timestamp) error {
			kvs[key] = value
			return nil
		}
		kc := NewCatalog(nil, snapshot).(*Catalog)
		ctx := context.Background()
		var collectionID int64 = 1
		var partitionID int64 = 2
		oldP := &model.Partition{PartitionID: partitionID, CollectionID: collectionID, State: pb.PartitionState_PartitionCreating}
		newP := &model.Partition{PartitionID: partitionID, CollectionID: collectionID, State: pb.PartitionState_PartitionCreated}
		err := kc.AlterPartition(ctx, testDb, oldP, newP, metastore.MODIFY, 0)
		assert.NoError(t, err)
		key := BuildPartitionKey(collectionID, partitionID)
		value, ok := kvs[key]
		assert.True(t, ok)
		var partPb pb.PartitionInfo
		err = proto.Unmarshal([]byte(value), &partPb)
		assert.NoError(t, err)
		got := model.UnmarshalPartitionModel(&partPb)
		assert.Equal(t, pb.PartitionState_PartitionCreated, got.State)
	})

	t.Run("modify, tenant id changed", func(t *testing.T) {
		kc := NewCatalog(nil, nil)
		ctx := context.Background()
		var collectionID int64 = 1
		oldP := &model.Partition{PartitionID: 1, CollectionID: collectionID, State: pb.PartitionState_PartitionCreating}
		newP := &model.Partition{PartitionID: 2, CollectionID: collectionID, State: pb.PartitionState_PartitionCreated}
		err := kc.AlterPartition(ctx, testDb, oldP, newP, metastore.MODIFY, 0)
		assert.Error(t, err)
	})
}

type mockSnapshotOpt func(ss *mocks.SnapShotKV)

func newMockSnapshot(t *testing.T, opts ...mockSnapshotOpt) *mocks.SnapShotKV {
	ss := mocks.NewSnapShotKV(t)
	for _, opt := range opts {
		opt(ss)
	}
	return ss
}

func withMockSave(saveErr error) mockSnapshotOpt {
	return func(ss *mocks.SnapShotKV) {
		ss.On(
			"Save",
			mock.Anything,
			mock.AnythingOfType("string"),
			mock.AnythingOfType("string"),
			mock.AnythingOfType("uint64")).
			Return(saveErr)
	}
}

func withMockMultiSave(multiSaveErr error) mockSnapshotOpt {
	return func(ss *mocks.SnapShotKV) {
		ss.On(
			"MultiSave",
			mock.Anything,
			mock.AnythingOfType("map[string]string"),
			mock.AnythingOfType("uint64")).
			Return(multiSaveErr)
	}
}

func withMockMultiSaveAndRemoveWithPrefix(err error) mockSnapshotOpt {
	return func(ss *mocks.SnapShotKV) {
		ss.On(
			"MultiSaveAndRemoveWithPrefix",
			mock.Anything,
			mock.AnythingOfType("map[string]string"),
			mock.AnythingOfType("[]string"),
			mock.AnythingOfType("uint64")).
			Return(err)
	}
}

func withMockMultiSaveAndRemove(err error) mockSnapshotOpt {
	return func(ss *mocks.SnapShotKV) {
		ss.On(
			"MultiSaveAndRemove",
			mock.Anything,
			mock.AnythingOfType("map[string]string"),
			mock.AnythingOfType("[]string"),
			mock.AnythingOfType("uint64")).
			Return(err)
	}
}

func TestCatalog_CreateCollection(t *testing.T) {
	t.Run("collection not creating", func(t *testing.T) {
		kc := NewCatalog(nil, nil)
		ctx := context.Background()
		coll := &model.Collection{State: pb.CollectionState_CollectionDropping}
		err := kc.CreateCollection(ctx, coll, 100)
		assert.Error(t, err)
	})

	t.Run("failed to save collection", func(t *testing.T) {
		mockSnapshot := newMockSnapshot(t, withMockSave(errors.New("error mock Save")))
		kc := NewCatalog(nil, mockSnapshot)
		ctx := context.Background()
		coll := &model.Collection{State: pb.CollectionState_CollectionCreated}
		err := kc.CreateCollection(ctx, coll, 100)
		assert.Error(t, err)
	})

	t.Run("succeed to save collection but failed to save other keys", func(t *testing.T) {
		mockSnapshot := newMockSnapshot(t, withMockSave(nil), withMockMultiSave(errors.New("error mock MultiSave")))
		kc := NewCatalog(nil, mockSnapshot)
		ctx := context.Background()
		coll := &model.Collection{
			Partitions: []*model.Partition{
				{PartitionName: "test"},
			},
			State: pb.CollectionState_CollectionCreated,
		}
		err := kc.CreateCollection(ctx, coll, 100)
		assert.Error(t, err)
	})

	t.Run("normal case", func(t *testing.T) {
		mockSnapshot := newMockSnapshot(t, withMockSave(nil), withMockMultiSave(nil))
		kc := NewCatalog(nil, mockSnapshot)
		ctx := context.Background()
		coll := &model.Collection{
			Partitions: []*model.Partition{
				{PartitionName: "test"},
			},
			State: pb.CollectionState_CollectionCreated,
		}
		err := kc.CreateCollection(ctx, coll, 100)
		assert.NoError(t, err)
	})

	t.Run("create collection with function and struct array field", func(t *testing.T) {
		mockSnapshot := newMockSnapshot(t, withMockSave(nil), withMockMultiSave(nil))
		kc := NewCatalog(nil, mockSnapshot)
		ctx := context.Background()
		coll := &model.Collection{
			Partitions: []*model.Partition{
				{PartitionName: "test"},
			},
			Fields: []*model.Field{
				{
					Name:     "text",
					DataType: schemapb.DataType_VarChar,
					TypeParams: []*commonpb.KeyValuePair{
						{
							Key:   "enable_analyzer",
							Value: "true",
						},
					},
				},
				{
					Name:     "sparse",
					DataType: schemapb.DataType_SparseFloatVector,
				},
			},
			StructArrayFields: []*model.StructArrayField{
				{
					Name: "test_struct",
					Fields: []*model.Field{
						{
							Name:        "sub_text",
							DataType:    schemapb.DataType_Array,
							ElementType: schemapb.DataType_VarChar,
						},
						{
							Name:        "sub_sparse",
							DataType:    schemapb.DataType_ArrayOfVector,
							ElementType: schemapb.DataType_SparseFloatVector,
						},
					},
				},
			},
			Functions: []*model.Function{
				{
					Name:             "test",
					Type:             schemapb.FunctionType_BM25,
					InputFieldNames:  []string{"text"},
					OutputFieldNames: []string{"sparse"},
				},
			},
			State: pb.CollectionState_CollectionCreated,
		}
		err := kc.CreateCollection(ctx, coll, 100)
		assert.NoError(t, err)
	})
}

func TestCatalog_DropCollection(t *testing.T) {
	t.Run("failed to remove", func(t *testing.T) {
		mockSnapshot := newMockSnapshot(t, withMockMultiSaveAndRemove(errors.New("error mock MultiSaveAndRemove")))
		kc := NewCatalog(nil, mockSnapshot)
		ctx := context.Background()
		coll := &model.Collection{
			Partitions: []*model.Partition{
				{PartitionName: "test"},
			},
			State: pb.CollectionState_CollectionDropping,
		}
		err := kc.DropCollection(ctx, coll, 100)
		assert.Error(t, err)
	})

	t.Run("succeed to remove first, but failed to remove twice", func(t *testing.T) {
		mockSnapshot := newMockSnapshot(t)
		removeOtherCalled := false
		removeCollectionCalled := false
		mockSnapshot.On(
			"MultiSaveAndRemove",
			mock.Anything,
			mock.AnythingOfType("map[string]string"),
			mock.AnythingOfType("[]string"),
			mock.AnythingOfType("uint64")).
			Return(func(context.Context, map[string]string, []string, typeutil.Timestamp) error {
				removeOtherCalled = true
				return nil
			}).Once()
		mockSnapshot.On(
			"MultiSaveAndRemove",
			mock.Anything,
			mock.AnythingOfType("map[string]string"),
			mock.AnythingOfType("[]string"),
			mock.AnythingOfType("uint64")).
			Return(func(context.Context, map[string]string, []string, typeutil.Timestamp) error {
				removeCollectionCalled = true
				return errors.New("error mock MultiSaveAndRemove")
			}).Once()
		kc := NewCatalog(nil, mockSnapshot)
		ctx := context.Background()
		coll := &model.Collection{
			Partitions: []*model.Partition{
				{PartitionName: "test"},
			},
			State: pb.CollectionState_CollectionDropping,
		}
		err := kc.DropCollection(ctx, coll, 100)
		assert.Error(t, err)
		assert.True(t, removeOtherCalled)
		assert.True(t, removeCollectionCalled)
	})

	t.Run("normal case", func(t *testing.T) {
		mockSnapshot := newMockSnapshot(t, withMockMultiSaveAndRemove(nil))
		kc := NewCatalog(nil, mockSnapshot)
		ctx := context.Background()
		coll := &model.Collection{
			Partitions: []*model.Partition{
				{PartitionName: "test"},
			},
			State: pb.CollectionState_CollectionDropping,
		}
		err := kc.DropCollection(ctx, coll, 100)
		assert.NoError(t, err)
	})

	t.Run("drop collection with function", func(t *testing.T) {
		mockSnapshot := newMockSnapshot(t, withMockMultiSaveAndRemove(nil))
		kc := NewCatalog(nil, mockSnapshot)
		ctx := context.Background()
		coll := &model.Collection{
			Partitions: []*model.Partition{
				{PartitionName: "test"},
			},
			Fields: []*model.Field{
				{
					Name:     "text",
					DataType: schemapb.DataType_VarChar,
					TypeParams: []*commonpb.KeyValuePair{
						{
							Key:   "enable_analyzer",
							Value: "true",
						},
					},
				},
				{
					Name:     "sparse",
					DataType: schemapb.DataType_SparseFloatVector,
				},
			},
			StructArrayFields: []*model.StructArrayField{
				{
					Name: "test_struct",
					Fields: []*model.Field{
						{
							Name:        "sub_text",
							DataType:    schemapb.DataType_Array,
							ElementType: schemapb.DataType_VarChar,
						},
						{
							Name:        "sub_sparse",
							DataType:    schemapb.DataType_ArrayOfVector,
							ElementType: schemapb.DataType_SparseFloatVector,
						},
					},
				},
			},
			Functions: []*model.Function{
				{
					Name:             "test",
					Type:             schemapb.FunctionType_BM25,
					InputFieldNames:  []string{"text"},
					OutputFieldNames: []string{"sparse"},
				},
			},
			State: pb.CollectionState_CollectionDropping,
		}
		err := kc.DropCollection(ctx, coll, 100)
		assert.NoError(t, err)
	})
}

func getUserInfoMetaString(username string) string {
	validInfo := &internalpb.CredentialInfo{Username: username, EncryptedPassword: "pwd" + username}
	validBytes, _ := json.Marshal(validInfo)
	return string(validBytes)
}

func TestRBAC_Credential(t *testing.T) {
	ctx := context.TODO()

	t.Run("test GetCredential", func(t *testing.T) {
		var (
			kvmock = mocks.NewTxnKV(t)
			c      = NewCatalog(kvmock, nil)

			loadFailName    = "invalid"
			loadFailKey     = fmt.Sprintf("%s/%s", CredentialPrefix, loadFailName)
			marshalFailName = "marshal"
			marshalFailkey  = fmt.Sprintf("%s/%s", CredentialPrefix, marshalFailName)
		)

		kvmock.EXPECT().Load(mock.Anything, loadFailKey).Return("", errors.New("Mock invalid load"))
		kvmock.EXPECT().Load(mock.Anything, marshalFailkey).Return("random", nil)
		kvmock.EXPECT().Load(mock.Anything, mock.Anything).Call.Return(
			func(ctx context.Context, key string) string {
				v, err := json.Marshal(&internalpb.CredentialInfo{EncryptedPassword: key})
				require.NoError(t, err)
				return string(v)
			}, nil)

		tests := []struct {
			description string
			isValid     bool

			user string

			expectedName     string
			expectedPassword string
		}{
			{"valid user1", true, "user1", "user1", "user1"},
			{"valid user2", true, "user2", "user2", "user2"},
			{"invalid user loadFail", false, loadFailName, "", ""},
			{"invalid user unmarshalFail", false, marshalFailName, "", ""},
		}

		for _, test := range tests {
			t.Run(test.description, func(t *testing.T) {
				cre, err := c.GetCredential(ctx, test.user)
				if test.isValid {
					assert.NoError(t, err)
					assert.Equal(t, test.expectedName, cre.Username)

					expectedPassword := fmt.Sprintf("%s/%s", CredentialPrefix, test.expectedPassword)
					assert.Equal(t, expectedPassword, cre.EncryptedPassword)
				} else {
					assert.Error(t, err)
				}
			})
		}
	})

	t.Run("test CreateCredential", func(t *testing.T) {
		var (
			kvmock      = mocks.NewTxnKV(t)
			c           = NewCatalog(kvmock, nil)
			invalidName = "invalid"
		)

		mockFailPath := fmt.Sprintf("%s/%s", CredentialPrefix, invalidName)
		kvmock.EXPECT().Save(mock.Anything, mockFailPath, mock.Anything).Return(errors.New("Mock invalid save"))
		kvmock.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil)

		tests := []struct {
			description string
			isValid     bool

			user     string
			password string
		}{
			{"valid user and password", true, "user1", "password"},
			{"valid user and password empty", true, "user2", ""},
			{"invalid user", false, invalidName, "password"},
		}

		for _, test := range tests {
			t.Run(test.description, func(t *testing.T) {
				err := c.AlterCredential(ctx, &model.Credential{
					Username:          test.user,
					EncryptedPassword: test.password,
				})

				if test.isValid {
					assert.NoError(t, err)
				} else {
					assert.Error(t, err)
				}

				err = c.AlterCredential(ctx, &model.Credential{
					Username:          test.user,
					EncryptedPassword: test.password,
				})

				if test.isValid {
					assert.NoError(t, err)
				} else {
					assert.Error(t, err)
				}
			})
		}
	})

	t.Run("test DropCredential", func(t *testing.T) {
		var (
			kvmock = mocks.NewTxnKV(t)
			c      = NewCatalog(kvmock, nil)

			validName              = "user1"
			validUserRoleKeyPrefix = funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, util.DefaultTenant, validName) + "/"

			dropFailName          = "drop-fail"
			dropUserRoleKeyPrefix = funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, util.DefaultTenant, dropFailName) + "/"
			getFailName           = "get-fail"
		)

		kvmock.EXPECT().MultiRemove(mock.Anything, []string{fmt.Sprintf("%s/%s", CredentialPrefix, dropFailName)}).Return(errors.New("Mock drop fail"))
		kvmock.EXPECT().MultiRemove(
			mock.Anything,
			[]string{
				fmt.Sprintf("%s/%s", CredentialPrefix, validName),
				validUserRoleKeyPrefix + "role1",
				validUserRoleKeyPrefix + "role2",
			},
		).Return(nil)
		kvmock.EXPECT().MultiRemove(mock.Anything, mock.Anything).Return(errors.New("Mock invalid multi remove"))

		kvmock.EXPECT().Load(mock.Anything, fmt.Sprintf("%s/%s", CredentialPrefix, getFailName)).Return("", errors.New("Mock invalid load"))
		kvmock.EXPECT().Load(mock.Anything, fmt.Sprintf("%s/%s", CredentialPrefix, validName)).Return(getUserInfoMetaString(validName), nil)
		kvmock.EXPECT().Load(mock.Anything, fmt.Sprintf("%s/%s", CredentialPrefix, dropFailName)).Return(getUserInfoMetaString(dropFailName), nil)

		kvmock.EXPECT().LoadWithPrefix(mock.Anything, validUserRoleKeyPrefix).Return(
			[]string{validUserRoleKeyPrefix + "role1", validUserRoleKeyPrefix + "role2"},
			[]string{"", ""},
			nil,
		)
		kvmock.EXPECT().LoadWithPrefix(mock.Anything, dropUserRoleKeyPrefix).Return([]string{}, []string{}, nil)

		tests := []struct {
			description string
			isValid     bool

			user string
		}{
			{"valid user1", true, validName},
			{"invalid user get-fail", false, getFailName},
			{"invalid user drop-fail", false, dropFailName},
		}

		for _, test := range tests {
			t.Run(test.description, func(t *testing.T) {
				err := c.DropCredential(ctx, test.user)
				if test.isValid {
					assert.NoError(t, err)
				} else {
					assert.Error(t, err)
				}
			})
		}
	})

	t.Run("test ListCredentials", func(t *testing.T) {
		var (
			kvmock = mocks.NewTxnKV(t)
			c      = NewCatalog(kvmock, nil)

			cmu   sync.RWMutex
			count = 0
		)

		// Return valid keys if count==0
		// return error if count!=0
		kvmock.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Call.Return(
			func(ctx context.Context, key string) []string {
				cmu.RLock()
				defer cmu.RUnlock()
				if count == 0 {
					return []string{
						fmt.Sprintf("%s/%s", CredentialPrefix, "user1"),
						fmt.Sprintf("%s/%s", CredentialPrefix, "user2"),
						fmt.Sprintf("%s/%s", CredentialPrefix, "user3"),
						"random",
					}
				}
				return nil
			},
			func(ctx context.Context, key string) []string {
				cmu.RLock()
				defer cmu.RUnlock()
				passwd, _ := json.Marshal(&model.Credential{EncryptedPassword: crypto.Base64Encode("passwd")})
				if count == 0 {
					return []string{
						string(passwd),
						string(passwd),
						string(passwd),
						string(passwd),
					}
				}
				return nil
			},
			func(ctx context.Context, key string) error {
				cmu.RLock()
				defer cmu.RUnlock()
				if count == 0 {
					return nil
				}
				return errors.New("Mock load with prefix")
			})

		tests := []struct {
			description string
			isValid     bool

			count       int
			expectedOut []string
		}{
			{"valid list", true, 0, []string{"user1", "user2", "user3"}},
			{"invalid list", false, 1, []string{}},
		}

		for _, test := range tests {
			t.Run(test.description, func(t *testing.T) {
				cmu.Lock()
				count = test.count
				cmu.Unlock()

				users, err := c.ListCredentials(ctx)
				if test.isValid {
					assert.NoError(t, err)
					assert.ElementsMatch(t, test.expectedOut, users)
				} else {
					assert.Error(t, err)
					assert.Empty(t, users)
				}
			})
		}
	})
}

func TestRBAC_Role(t *testing.T) {
	ctx := context.TODO()
	tenant := "default"

	t.Run("test remove", func(t *testing.T) {
		var (
			kvmock = mocks.NewTxnKV(t)
			c      = NewCatalog(kvmock, nil).(*Catalog)

			notExistKey = "not-exist"
			errorKey    = "error"
			otherError  = errors.New("mock load error")
		)

		kvmock.EXPECT().Load(mock.Anything, notExistKey).Return("", merr.WrapErrIoKeyNotFound(notExistKey)).Once()
		kvmock.EXPECT().Load(mock.Anything, errorKey).Return("", otherError).Once()
		kvmock.EXPECT().Load(mock.Anything, mock.Anything).Return("", nil).Once()
		kvmock.EXPECT().Remove(mock.Anything, mock.Anything).Call.Return(nil).Once()
		tests := []struct {
			description string

			isValid bool
			key     string

			expectedError error
			ignorable     bool
		}{
			{"error key not exists, ignorable", false, notExistKey, nil, true},
			{"error other error", false, errorKey, otherError, false},
			{"no error", true, "key1", nil, false},
		}
		for _, test := range tests {
			t.Run(test.description, func(t *testing.T) {
				err := c.remove(ctx, test.key)
				if test.isValid {
					assert.NoError(t, err)
				} else {
					assert.Error(t, err)
				}

				if test.ignorable {
					_, ok := err.(*common.IgnorableError)
					assert.True(t, ok)
				}
			})
		}
	})
	t.Run("test CreateRole", func(t *testing.T) {
		var (
			kvmock = mocks.NewTxnKV(t)
			c      = NewCatalog(kvmock, nil)
		)

		kvmock.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Call.Return(nil).Once()
		err := c.CreateRole(ctx, tenant, &milvuspb.RoleEntity{
			Name: "role1",
		})
		assert.NoError(t, err)
	})
	t.Run("test DropRole", func(t *testing.T) {
		var (
			kvmock = mocks.NewTxnKV(t)
			c      = NewCatalog(kvmock, nil)

			validName   = "role1"
			errorName   = "error"
			getFailName = "get-fail"
		)

		kvmock.EXPECT().MultiRemove(mock.Anything, []string{funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, errorName)}).Return(errors.New("remove error"))
		kvmock.EXPECT().MultiRemove(mock.Anything, []string{
			funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, validName),
			funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, fmt.Sprintf("%s/%s", "user1", validName)),
			funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, fmt.Sprintf("%s/%s", "user2", validName)),
		}).Return(nil)
		kvmock.EXPECT().MultiRemove(mock.Anything, mock.Anything).Return(errors.New("mock multi remove error"))

		getRoleMappingKey := func(username, rolename string) string {
			return funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, fmt.Sprintf("%s/%s", username, rolename))
		}

		kvmock.EXPECT().LoadWithPrefix(mock.Anything, funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, "")).Return(
			[]string{getRoleMappingKey("user1", validName), getRoleMappingKey("user2", validName), getRoleMappingKey("user3", "role3")},
			[]string{},
			nil,
		)

		kvmock.EXPECT().Load(mock.Anything, funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, getFailName)).Return("", errors.New("mock load error"))
		kvmock.EXPECT().Load(mock.Anything, funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, validName)).Return("", nil)
		kvmock.EXPECT().Load(mock.Anything, funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, errorName)).Return("", nil)

		tests := []struct {
			description string
			isValid     bool
			role        string
		}{
			{"valid role role1", true, validName},
			{"fail to get role info", false, getFailName},
			{"invalid role error", false, errorName},
		}

		for _, test := range tests {
			t.Run(test.description, func(t *testing.T) {
				err := c.DropRole(ctx, tenant, test.role)
				if test.isValid {
					assert.NoError(t, err)
				} else {
					assert.Error(t, err)
				}
			})
		}
	})
	t.Run("test AlterUserRole", func(t *testing.T) {
		var (
			kvmock = mocks.NewTxnKV(t)
			c      = NewCatalog(kvmock, nil)

			user = "default-user"
		)
		kvmock.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(nil)
		kvmock.EXPECT().Remove(mock.Anything, mock.Anything).Return(nil)
		err := c.AlterUserRole(ctx, tenant, &milvuspb.UserEntity{Name: user}, &milvuspb.RoleEntity{
			Name: "role1",
		}, milvuspb.OperateUserRoleType_RemoveUserFromRole)
		require.NoError(t, err)
		err = c.AlterUserRole(ctx, tenant, &milvuspb.UserEntity{Name: user}, &milvuspb.RoleEntity{
			Name: "role1",
		}, milvuspb.OperateUserRoleType_AddUserToRole)
		require.NoError(t, err)
	})

	t.Run("test ListRole", func(t *testing.T) {
		var loadWithPrefixReturn atomic.Bool

		t.Run("test entity!=nil", func(t *testing.T) {
			var (
				kvmock = mocks.NewTxnKV(t)
				c      = NewCatalog(kvmock, nil)

				errorLoad     = "error"
				errorLoadPath = funcutil.HandleTenantForEtcdKey(RolePrefix, tenant, errorLoad)
			)

			kvmock.EXPECT().Load(mock.Anything, errorLoadPath).Call.Return("", errors.New("mock load error"))
			kvmock.EXPECT().Load(mock.Anything, mock.Anything).Call.Return("", nil)

			// Return valid keys if loadWithPrefixReturn == True
			// return error if loadWithPrefixReturn == False
			kvmock.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Call.Return(
				func(ctx context.Context, key string) []string {
					if loadWithPrefixReturn.Load() {
						return []string{
							fmt.Sprintf("%s/%s/%s", RoleMappingPrefix, tenant, "user1/role1"),
							fmt.Sprintf("%s/%s/%s", RoleMappingPrefix, tenant, "user2/role2"),
							fmt.Sprintf("%s/%s/%s", RoleMappingPrefix, tenant, "user3/role3"),
							"random",
						}
					}
					return nil
				},
				nil,
				func(ctx context.Context, key string) error {
					if loadWithPrefixReturn.Load() {
						return nil
					}
					return fmt.Errorf("Mock load with prefix wrong, key=%s", key)
				})

			tests := []struct {
				description string
				isValid     bool

				includeUserInfo      bool
				loadWithPrefixReturn bool

				entity *milvuspb.RoleEntity
			}{
				{"loadWithPrefix error", false, true, false, &milvuspb.RoleEntity{Name: "role1"}},
				{"entity empty name", false, true, true, &milvuspb.RoleEntity{Name: ""}},
				{"entity load error", false, true, true, &milvuspb.RoleEntity{Name: errorLoad}},
				{"entity role1", true, true, true, &milvuspb.RoleEntity{Name: "role1"}},
				{"entity empty name/includeUserInfo=False", false, false, true, &milvuspb.RoleEntity{Name: ""}},
				{"entity load error/includeUserInfo=False", false, false, true, &milvuspb.RoleEntity{Name: errorLoad}},
				{"entity role1/includeUserInfo=False", true, false, true, &milvuspb.RoleEntity{Name: "role1"}},
			}

			for _, test := range tests {
				t.Run(test.description, func(t *testing.T) {
					loadWithPrefixReturn.Store(test.loadWithPrefixReturn)
					res, err := c.ListRole(ctx, tenant, test.entity, test.includeUserInfo)
					if test.isValid {
						assert.NoError(t, err)

						assert.Equal(t, 1, len(res))
						assert.Equal(t, test.entity.Name, res[0].GetRole().GetName())
						if test.includeUserInfo {
							assert.Equal(t, 1, len(res[0].GetUsers()))
							assert.Equal(t, "user1", res[0].GetUsers()[0].GetName())
						} else {
							for _, r := range res {
								assert.Empty(t, r.GetUsers())
							}
						}
					} else {
						assert.Error(t, err)
					}
				})
			}
		})

		t.Run("test entity is nil", func(t *testing.T) {
			var (
				kvmock = mocks.NewTxnKV(t)
				c      = NewCatalog(kvmock, nil)
			)

			// Return valid keys if loadWithPrefixReturn == True
			// return error if loadWithPrefixReturn == False
			// Mocking the return of kv_catalog.go:L699
			kvmock.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Call.Return(
				func(ctx context.Context, key string) []string {
					if loadWithPrefixReturn.Load() {
						return []string{
							fmt.Sprintf("%s/%s/%s", RolePrefix, tenant, "role1"),
							fmt.Sprintf("%s/%s/%s", RolePrefix, tenant, "role2"),
							fmt.Sprintf("%s/%s/%s", RolePrefix, tenant, "role3"),
							"random",
						}
					}
					return nil
				},
				nil,
				func(ctx context.Context, key string) error {
					if loadWithPrefixReturn.Load() {
						return nil
					}
					return fmt.Errorf("Mock load with prefix wrong, key=%s", key)
				})

			tests := []struct {
				description string
				isValid     bool

				loadWithPrefixReturn bool
			}{
				{"entity nil/loadWithPrefix success", true, true},
				{"entity nil/loadWithPrefix fail", false, false},
			}

			for _, test := range tests {
				t.Run(test.description, func(t *testing.T) {
					loadWithPrefixReturn.Store(test.loadWithPrefixReturn)
					res, err := c.ListRole(ctx, tenant, nil, false)
					if test.isValid {
						assert.NoError(t, err)
						assert.Equal(t, 3, len(res))

						for _, r := range res {
							assert.Empty(t, r.GetUsers())
						}
					} else {
						assert.Error(t, err)
					}
				})
			}
		})
	})
	t.Run("test ListUser", func(t *testing.T) {
		var (
			kvmock = mocks.NewTxnKV(t)
			c      = NewCatalog(kvmock, nil).(*Catalog)

			invalidUser    = "invalid-user"
			invalidUserKey = funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, invalidUser) + "/"
		)
		// returns error for invalidUserKey
		kvmock.EXPECT().LoadWithPrefix(mock.Anything, invalidUserKey).Call.Return(
			nil, nil, errors.New("Mock load with prefix wrong"))

		// Returns keys for RoleMappingPrefix/tenant/user1/ (with trailing slash after the fix)
		user1Key := funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, "user1") + "/"
		kvmock.EXPECT().LoadWithPrefix(mock.Anything, user1Key).Call.Return(
			func(ctx context.Context, key string) []string {
				return []string{
					user1Key + "role1",
					user1Key + "role2",
					user1Key + "role3/error",
				}
			}, nil, nil)

		// Returns keys for CredentialPrefix
		var loadCredentialPrefixReturn atomic.Bool
		kvmock.EXPECT().LoadWithPrefix(mock.Anything, CredentialPrefix).Call.Return(
			func(ctx context.Context, key string) []string {
				if loadCredentialPrefixReturn.Load() {
					return []string{
						fmt.Sprintf("%s/%s/%s", CredentialPrefix, UserSubPrefix, "user1"),
					}
				}
				return nil
			},
			func(ctx context.Context, key string) []string {
				if loadCredentialPrefixReturn.Load() {
					passwd, _ := json.Marshal(&model.Credential{EncryptedPassword: crypto.Base64Encode("passwd")})
					return []string{string(passwd)}
				}
				return nil
			},
			func(ctx context.Context, key string) error {
				if loadCredentialPrefixReturn.Load() {
					return nil
				}

				return fmt.Errorf("mock load with prefix error, key=%s", key)
			})

		t.Run("test getUserResult", func(t *testing.T) {
			tests := []struct {
				description string
				isValid     bool

				user            string
				includeRoleInfo bool
			}{
				{"valid user1 not include RoleInfo", true, "user1", false},
				{"valid user1 include RoleInfo", true, "user1", true},
				{"invalid user not include RoleInfo", true, invalidUser, false},
				{"invalid user include RoleInfo", false, invalidUser, true},
			}

			for _, test := range tests {
				t.Run(test.description, func(t *testing.T) {
					res, err := c.getUserResult(ctx, tenant, test.user, test.includeRoleInfo)

					assert.Equal(t, test.user, res.GetUser().GetName())

					if test.isValid {
						assert.NoError(t, err)

						if test.includeRoleInfo {
							assert.Equal(t, 2, len(res.GetRoles()))
							assert.Equal(t, "role1", res.GetRoles()[0].GetName())
							assert.Equal(t, "role2", res.GetRoles()[1].GetName())
						} else {
							assert.Equal(t, 0, len(res.GetRoles()))
						}
					} else {
						assert.Error(t, err)
					}
				})
			}
		})

		t.Run("test ListUser", func(t *testing.T) {
			var (
				invalidUserLoad    = "invalid-user-load"
				invalidUserLoadKey = fmt.Sprintf("%s/%s", CredentialPrefix, invalidUserLoad)
			)
			// Returns error for invalidUserLoadKey
			kvmock.EXPECT().Load(mock.Anything, invalidUserLoadKey).Call.Return("", errors.New("Mock load wrong"))
			kvmock.EXPECT().Load(mock.Anything, mock.Anything).Call.Return(
				func(ctx context.Context, key string) string {
					v, err := json.Marshal(&internalpb.CredentialInfo{EncryptedPassword: key})
					require.NoError(t, err)
					return string(v)
				}, nil)
			tests := []struct {
				isValid                    bool
				loadCredentialPrefixReturn bool

				entity          *milvuspb.UserEntity
				includeRoleInfo bool

				description string
			}{
				{true, true, nil, true, "nil entity include RoleInfo"},
				{true, true, nil, false, "nil entity not include RoleInfo"},
				{false, false, nil, false, "nil entity ListCredentials error"},
				{false, false, nil, true, "nil entity ListCredentials error"},
				{false, true, &milvuspb.UserEntity{Name: ""}, false, "empty entity Name"},
				{false, true, &milvuspb.UserEntity{Name: invalidUserLoad}, false, "invalid entity Name"},
				{true, true, &milvuspb.UserEntity{Name: "user1"}, false, "valid entity user1 not include RoleInfo"},
				{true, true, &milvuspb.UserEntity{Name: "user1"}, true, "valid entity user1 include RoleInfo"},
				{false, true, &milvuspb.UserEntity{Name: invalidUser}, true, "invalid entity invalidUser include RoleInfo"},
			}

			for _, test := range tests {
				t.Run(test.description, func(t *testing.T) {
					loadCredentialPrefixReturn.Store(test.loadCredentialPrefixReturn)
					res, err := c.ListUser(ctx, tenant, test.entity, test.includeRoleInfo)

					if test.isValid {
						assert.NoError(t, err)
						assert.Equal(t, 1, len(res))
						assert.Equal(t, "user1", res[0].GetUser().GetName())
						if test.includeRoleInfo {
							assert.Equal(t, 2, len(res[0].GetRoles()))
						} else {
							assert.Equal(t, 0, len(res[0].GetRoles()))
						}
					} else {
						assert.Error(t, err)
						assert.Empty(t, res)
					}
				})
			}
		})
	})
	t.Run("test ListUserRole", func(t *testing.T) {
		var (
			loadWithPrefixReturn atomic.Bool
			kvmock               = mocks.NewTxnKV(t)
			c                    = NewCatalog(kvmock, nil)
		)

		// Return valid keys if loadWithPrefixReturn == True
		// return error if loadWithPrefixReturn == False
		// Mocking the return of kv_catalog.go:ListUserRole:L982
		kvmock.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Call.Return(
			func(ctx context.Context, key string) []string {
				if loadWithPrefixReturn.Load() {
					return []string{
						fmt.Sprintf("%s/%s/%s", RoleMappingPrefix, tenant, "user1/role1"),
						fmt.Sprintf("%s/%s/%s", RoleMappingPrefix, tenant, "user1/role2"),
						fmt.Sprintf("%s/%s/%s", RoleMappingPrefix, tenant, "user3/role3"),
						"random",
					}
				}
				return nil
			},
			nil,
			func(ctx context.Context, key string) error {
				if loadWithPrefixReturn.Load() {
					return nil
				}
				return fmt.Errorf("Mock load with prefix wrong, key=%s", key)
			})

		tests := []struct {
			isValid              bool
			loadWithPrefixReturn bool

			description string
		}{
			{true, true, "valid loadWithPrefix"},
			{false, false, "error loadWithPrefix"},
		}

		for _, test := range tests {
			t.Run(test.description, func(t *testing.T) {
				loadWithPrefixReturn.Store(test.loadWithPrefixReturn)
				res, err := c.ListUserRole(ctx, tenant)
				if test.isValid {
					assert.NoError(t, err)
					assert.Equal(t, 3, len(res))
					assert.ElementsMatch(t, []string{"user1/role1", "user1/role2", "user3/role3"}, res)
				} else {
					assert.Error(t, err)
				}
			})
		}
	})
}

func TestRBAC_Grant(t *testing.T) {
	var (
		tenant = "default"
		ctx    = context.TODO()

		object  = "Collection"
		objName = "insert"

		validRole       = "role1"
		invalidRole     = "role2"
		keyNotExistRole = "role3"
		errorSaveRole   = "role100"

		validUser   = "user1"
		invalidUser = "user2"

		validPrivilege        = "write"
		invalidPrivilege      = "wal"
		keyNotExistPrivilege  = "read"
		keyNotExistPrivilege2 = "read2"
	)
	t.Run("test AlterGrant", func(t *testing.T) {
		var (
			kvmock = mocks.NewTxnKV(t)
			c      = NewCatalog(kvmock, nil)
		)

		validRoleKey := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, fmt.Sprintf("%s/%s/%s", validRole, object, objName))
		validRoleValue := crypto.MD5(validRoleKey)

		invalidRoleKey := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, fmt.Sprintf("%s/%s/%s", invalidRole, object, objName))
		invalidRoleKeyWithDb := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, fmt.Sprintf("%s/%s/%s", invalidRole, object, funcutil.CombineObjectName(util.DefaultDBName, objName)))

		keyNotExistRoleKey := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, fmt.Sprintf("%s/%s/%s", keyNotExistRole, object, objName))
		keyNotExistRoleKeyWithDb := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, fmt.Sprintf("%s/%s/%s", keyNotExistRole, object, funcutil.CombineObjectName(util.DefaultDBName, objName)))
		keyNotExistRoleValueWithDb := crypto.MD5(keyNotExistRoleKeyWithDb)

		errorSaveRoleKey := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, fmt.Sprintf("%s/%s/%s", errorSaveRole, object, objName))
		errorSaveRoleKeyWithDb := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, fmt.Sprintf("%s/%s/%s", errorSaveRole, object, funcutil.CombineObjectName(util.DefaultDBName, objName)))

		// Mock return in kv_catalog.go:AlterGrant:L815
		kvmock.EXPECT().Load(mock.Anything, validRoleKey).Call.
			Return(func(ctx context.Context, key string) string { return validRoleValue }, nil)

		kvmock.EXPECT().Load(mock.Anything, invalidRoleKey).Call.
			Return("", func(ctx context.Context, key string) error {
				return fmt.Errorf("mock load error, key=%s", key)
			})
		kvmock.EXPECT().Load(mock.Anything, invalidRoleKeyWithDb).Call.
			Return("", func(ctx context.Context, key string) error {
				return fmt.Errorf("mock load error, key=%s", key)
			})
		kvmock.EXPECT().Load(mock.Anything, keyNotExistRoleKey).Call.
			Return("", func(ctx context.Context, key string) error {
				return merr.WrapErrIoKeyNotFound(key)
			})
		kvmock.EXPECT().Load(mock.Anything, keyNotExistRoleKeyWithDb).Call.
			Return("", func(ctx context.Context, key string) error {
				return merr.WrapErrIoKeyNotFound(key)
			})
		kvmock.EXPECT().Load(mock.Anything, errorSaveRoleKey).Call.
			Return("", func(ctx context.Context, key string) error {
				return merr.WrapErrIoKeyNotFound(key)
			})
		kvmock.EXPECT().Load(mock.Anything, errorSaveRoleKeyWithDb).Call.
			Return("", func(ctx context.Context, key string) error {
				return merr.WrapErrIoKeyNotFound(key)
			})
		kvmock.EXPECT().Save(mock.Anything, keyNotExistRoleKeyWithDb, mock.Anything).Return(nil)
		kvmock.EXPECT().Save(mock.Anything, errorSaveRoleKeyWithDb, mock.Anything).Return(errors.New("mock save error role"))

		validPrivilegeKey := funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, fmt.Sprintf("%s/%s", validRoleValue, validPrivilege))
		invalidPrivilegeKey := funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, fmt.Sprintf("%s/%s", validRoleValue, invalidPrivilege))
		keyNotExistPrivilegeKey := funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, fmt.Sprintf("%s/%s", validRoleValue, keyNotExistPrivilege))
		keyNotExistPrivilegeKey2WithDb := funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, fmt.Sprintf("%s/%s", keyNotExistRoleValueWithDb, keyNotExistPrivilege2))

		// Mock return in kv_catalog.go:AlterGrant:L838
		kvmock.EXPECT().Load(mock.Anything, validPrivilegeKey).Call.Return("", nil)
		kvmock.EXPECT().Load(mock.Anything, invalidPrivilegeKey).Call.
			Return("", func(ctx context.Context, key string) error {
				return fmt.Errorf("mock load error, key=%s", key)
			})
		kvmock.EXPECT().Load(mock.Anything, keyNotExistPrivilegeKey).Call.
			Return("", func(ctx context.Context, key string) error {
				return merr.WrapErrIoKeyNotFound(key)
			})
		kvmock.EXPECT().Load(mock.Anything, keyNotExistPrivilegeKey2WithDb).Call.
			Return("", func(ctx context.Context, key string) error {
				return merr.WrapErrIoKeyNotFound(key)
			})
		kvmock.EXPECT().Load(mock.Anything, mock.Anything).Call.Return("", nil)

		t.Run("test Grant", func(t *testing.T) {
			kvmock.EXPECT().Save(mock.Anything, mock.Anything, validUser).Return(nil)
			kvmock.EXPECT().Save(mock.Anything, mock.Anything, invalidUser).Return(errors.New("mock save invalid user"))

			tests := []struct {
				isValid bool

				userName      string
				roleName      string
				privilegeName string

				ignorable   bool
				description string
			}{
				// exist role
				{false, validUser, validRole, invalidPrivilege, false, "grant exist Role with error Privilege"},
				{false, validUser, validRole, validPrivilege, true, "grant exist Role with exist Privilege, ignorable"},
				{false, invalidUser, validRole, keyNotExistPrivilege, false, "grant exist Role with not exist Privilege with invalid user"},
				{true, validUser, validRole, keyNotExistPrivilege, true, "grant exist Role with not exist Privilege with valid user"},
				// error role
				{false, validUser, invalidRole, invalidPrivilege, false, "grant invalid role with invalid privilege"},
				{false, validUser, invalidRole, validPrivilege, false, "grant invalid role with valid privilege"},
				{false, validUser, invalidRole, keyNotExistPrivilege, false, "grant invalid role with not exist privilege"},
				{false, validUser, errorSaveRole, keyNotExistPrivilege, false, "grant error role with not exist privilege"},
				// not exist role
				{false, validUser, keyNotExistRole, validPrivilege, true, "grant not exist role with exist privilege"},
				{true, validUser, keyNotExistRole, keyNotExistPrivilege2, false, "grant not exist role with not exist privilege"},
			}

			for _, test := range tests {
				t.Run(test.description, func(t *testing.T) {
					err := c.AlterGrant(ctx, tenant, &milvuspb.GrantEntity{
						Role:       &milvuspb.RoleEntity{Name: test.roleName},
						Object:     &milvuspb.ObjectEntity{Name: object},
						ObjectName: objName,
						DbName:     util.DefaultDBName,
						Grantor: &milvuspb.GrantorEntity{
							User:      &milvuspb.UserEntity{Name: test.userName},
							Privilege: &milvuspb.PrivilegeEntity{Name: test.privilegeName},
						},
					}, milvuspb.OperatePrivilegeType_Grant)

					if test.isValid {
						assert.NoError(t, err)
					} else {
						assert.Error(t, err)
						_, ok := err.(*common.IgnorableError)
						if test.ignorable {
							assert.True(t, ok)
						} else {
							assert.False(t, ok)
						}
					}
				})
			}
		})

		t.Run("test Revoke", func(t *testing.T) {
			invalidPrivilegeRemove := "p-remove"
			invalidPrivilegeRemoveKey := funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, fmt.Sprintf("%s/%s", validRoleValue, invalidPrivilegeRemove))

			kvmock.EXPECT().Load(mock.Anything, invalidPrivilegeRemoveKey).Call.Return("", nil)
			kvmock.EXPECT().Remove(mock.Anything, invalidPrivilegeRemoveKey).Return(errors.New("mock remove error"))
			kvmock.EXPECT().Remove(mock.Anything, mock.Anything).Return(nil)
			tests := []struct {
				isValid bool

				userName      string
				roleName      string
				privilegeName string

				ignorable   bool
				description string
			}{
				// invalid role
				{false, validUser, invalidRole, validPrivilege, false, "invalid role"},
				{false, validUser, invalidRole, invalidPrivilege, false, "invalid role, invalid privilege"},
				{false, validUser, invalidRole, keyNotExistPrivilege, false, "invalid role, not exist privilege"},
				// not exist role
				{false, validUser, keyNotExistRole, validPrivilege, true, "not exist role with exist privilege"},
				// exist role
				{false, validUser, validRole, invalidPrivilege, false, "exist role with invalid privilege"},
				{false, validUser, validRole, keyNotExistPrivilege, true, "exist role with not exist privilege"},
				{true, validUser, validRole, validPrivilege, false, "exist role with exist privilege success to remove"},
				{false, validUser, validRole, invalidPrivilegeRemove, false, "exist role with exist privilege fail to remove"},
			}

			for _, test := range tests {
				t.Run(test.description, func(t *testing.T) {
					err := c.AlterGrant(ctx, tenant, &milvuspb.GrantEntity{
						Role:       &milvuspb.RoleEntity{Name: test.roleName},
						Object:     &milvuspb.ObjectEntity{Name: object},
						ObjectName: objName,
						DbName:     util.DefaultDBName,
						Grantor: &milvuspb.GrantorEntity{
							User:      &milvuspb.UserEntity{Name: test.userName},
							Privilege: &milvuspb.PrivilegeEntity{Name: test.privilegeName},
						},
					}, milvuspb.OperatePrivilegeType_Revoke)

					if test.isValid {
						assert.NoError(t, err)
					} else {
						assert.Error(t, err)
						_, ok := err.(*common.IgnorableError)
						if test.ignorable {
							assert.True(t, ok)
						} else {
							assert.False(t, ok)
						}
					}
				})
			}
		})
	})
	t.Run("test DeleteGrant", func(t *testing.T) {
		var (
			kvmock = mocks.NewTxnKV(t)
			c      = NewCatalog(kvmock, nil)

			errorRole           = "error-role"
			errorRolePrefix     = funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, errorRole+"/")
			loadErrorRole       = "load-error-role"
			loadErrorRolePrefix = funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, loadErrorRole+"/")
			granteeID           = "123456"
			granteePrefix       = funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, granteeID+"/")
		)

		kvmock.EXPECT().LoadWithPrefix(mock.Anything, loadErrorRolePrefix).Call.Return(nil, nil, errors.New("mock loadWithPrefix error"))
		kvmock.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Call.Return(nil, []string{granteeID}, nil)
		kvmock.EXPECT().MultiSaveAndRemoveWithPrefix(mock.Anything, mock.Anything, []string{errorRolePrefix, granteePrefix}, mock.Anything).Call.Return(errors.New("mock removeWithPrefix error"))
		kvmock.EXPECT().MultiSaveAndRemoveWithPrefix(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Call.Return(nil)

		tests := []struct {
			isValid bool
			role    string

			description string
		}{
			{true, "role1", "valid role1"},
			{false, errorRole, "invalid errorRole"},
			{false, loadErrorRole, "invalid load errorRole"},
		}

		for _, test := range tests {
			t.Run(test.description, func(t *testing.T) {
				err := c.DeleteGrant(ctx, tenant, &milvuspb.RoleEntity{Name: test.role})
				if test.isValid {
					assert.NoError(t, err)
				} else {
					assert.Error(t, err)
				}
			})
		}
	})
	t.Run("test ListGrant", func(t *testing.T) {
		var (
			kvmock = mocks.NewTxnKV(t)
			c      = NewCatalog(kvmock, nil)
		)

		// Mock Load in kv_catalog.go:L901
		validGranteeKey := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant,
			fmt.Sprintf("%s/%s/%s", "role1", "obj1", "obj_name1"))
		kvmock.EXPECT().Load(mock.Anything, validGranteeKey).Call.
			Return(func(ctx context.Context, key string) string { return crypto.MD5(key) }, nil)
		validGranteeKey2 := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant,
			fmt.Sprintf("%s/%s/%s", "role1", "obj2", "foo.obj_name2"))
		kvmock.EXPECT().Load(mock.Anything, validGranteeKey2).Call.
			Return(func(ctx context.Context, key string) string { return crypto.MD5(key) }, nil)
		validGranteeKey3 := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant,
			fmt.Sprintf("%s/%s/%s", "role1", "obj3", "*.obj_name3"))
		kvmock.EXPECT().Load(mock.Anything, validGranteeKey3).Call.
			Return(func(ctx context.Context, key string) string { return crypto.MD5(key) }, nil)

		kvmock.EXPECT().Load(mock.Anything, mock.Anything).Call.
			Return("", errors.New("mock Load error"))

		invalidRoleKey := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, invalidRole) + "/"
		kvmock.EXPECT().LoadWithPrefix(mock.Anything, invalidRoleKey).Call.Return(nil, nil, errors.New("mock loadWithPrefix error"))
		kvmock.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Call.Return(
			func(ctx context.Context, key string) []string {
				// Mock kv_catalog.go:ListGrant - key now ends with "/" after the fix
				if strings.Contains(key, GranteeIDPrefix) {
					return []string{
						key + "PrivilegeLoad",
						key + "PrivilegeRelease",
					}
				}

				// Mock kv_catalog.go:ListGrant - key now ends with "/" after the fix
				return []string{
					key + "obj1/obj_name1",
					key + "obj2/foo.obj_name2",
					key + "obj3/*.obj_name3",
				}
			},
			func(ctx context.Context, key string) []string {
				if strings.Contains(key, GranteeIDPrefix) {
					return []string{"user1", "user2"}
				}
				return []string{
					crypto.MD5(key + "obj1/obj_name1"),
					crypto.MD5(key + "obj2/foo.obj_name2"),
					crypto.MD5(key + "obj3/*.obj_name3"),
				}
			},
			nil,
		)

		tests := []struct {
			isValid bool

			entity      *milvuspb.GrantEntity
			description string
			count       int
		}{
			{true, &milvuspb.GrantEntity{Role: &milvuspb.RoleEntity{Name: "role1"}}, "valid role role1 with empty entity", 4},
			{false, &milvuspb.GrantEntity{Role: &milvuspb.RoleEntity{Name: invalidRole}}, "invalid role with empty entity", 0},
			{false, &milvuspb.GrantEntity{
				Object:     &milvuspb.ObjectEntity{Name: "random"},
				ObjectName: "random2",
				Role:       &milvuspb.RoleEntity{Name: "role1"},
			}, "valid role with not exist entity", 0},
			{true, &milvuspb.GrantEntity{
				Object:     &milvuspb.ObjectEntity{Name: "obj1"},
				ObjectName: "obj_name1",
				Role:       &milvuspb.RoleEntity{Name: "role1"},
			}, "valid role with valid entity", 2},
			{true, &milvuspb.GrantEntity{
				Object:     &milvuspb.ObjectEntity{Name: "obj2"},
				ObjectName: "obj_name2",
				DbName:     "foo",
				Role:       &milvuspb.RoleEntity{Name: "role1"},
			}, "valid role and dbName with valid entity", 2},
			{false, &milvuspb.GrantEntity{
				Object:     &milvuspb.ObjectEntity{Name: "obj2"},
				ObjectName: "obj_name2",
				DbName:     "foo2",
				Role:       &milvuspb.RoleEntity{Name: "role1"},
			}, "valid role and invalid dbName with valid entity", 0},
			{false, &milvuspb.GrantEntity{
				Object:     &milvuspb.ObjectEntity{Name: "obj3"},
				ObjectName: "obj_name3",
				DbName:     "default",
				Role:       &milvuspb.RoleEntity{Name: "role1"},
			}, "valid role and dbName with default db", 2},
			{true, &milvuspb.GrantEntity{
				DbName: "default",
				Role:   &milvuspb.RoleEntity{Name: "role1"},
			}, "valid role and default dbName without object", 4},
			{true, &milvuspb.GrantEntity{
				DbName: "*",
				Role:   &milvuspb.RoleEntity{Name: "role1"},
			}, "valid role and any dbName without object", 6},
		}

		for _, test := range tests {
			t.Run(test.description, func(t *testing.T) {
				if test.entity.DbName == "" {
					test.entity.DbName = util.DefaultDBName
				}
				grants, err := c.ListGrant(ctx, tenant, test.entity)
				if test.isValid {
					assert.NoError(t, err)
					for _, g := range grants {
						assert.Equal(t, test.entity.GetRole().GetName(), g.GetRole().GetName())
					}
				} else {
					assert.Error(t, err)
				}
				assert.Equal(t, test.count, len(grants))
			})
		}
	})
	t.Run("test ListPolicy", func(t *testing.T) {
		var (
			kvmock = mocks.NewTxnKV(t)
			c      = NewCatalog(kvmock, nil)

			firstLoadWithPrefixReturn  atomic.Bool
			secondLoadWithPrefixReturn atomic.Bool
		)

		grant := func(role, obj, objName, privilege, dbName string) *milvuspb.GrantEntity {
			privilegeName := util.PrivilegeNameForAPI(privilege)
			if privilege == util.AnyWord {
				privilegeName = util.AnyWord
			}
			return &milvuspb.GrantEntity{
				Role:       &milvuspb.RoleEntity{Name: role},
				Object:     &milvuspb.ObjectEntity{Name: obj},
				ObjectName: objName,
				DbName:     dbName,
				Grantor: &milvuspb.GrantorEntity{
					Privilege: &milvuspb.PrivilegeEntity{Name: privilegeName},
				},
			}
		}

		kvmock.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Call.Return(
			func(ctx context.Context, key string) []string {
				contains := strings.Contains(key, GranteeIDPrefix)
				if contains {
					if secondLoadWithPrefixReturn.Load() {
						return []string{
							fmt.Sprintf("%s/%s", key, "PrivilegeLoad"),
							fmt.Sprintf("%s/%s", key, "PrivilegeRelease"),
							fmt.Sprintf("%s/%s", key, "random/a/b/c"),
							fmt.Sprintf("%s/%s", key, util.AnyWord),
						}
					}
					return nil
				}

				if firstLoadWithPrefixReturn.Load() {
					return []string{
						fmt.Sprintf("%s/%s", key, "role1/obj1/obj_name1"),
						fmt.Sprintf("%s/%s", key, "role2/obj2/obj_name2"),
						"random",
					}
				}
				return nil
			},

			func(ctx context.Context, key string) []string {
				if firstLoadWithPrefixReturn.Load() {
					return []string{
						crypto.MD5(fmt.Sprintf("%s/%s", key, "obj1/obj_name1")),
						crypto.MD5(fmt.Sprintf("%s/%s", key, "obj2/obj_name2")),
					}
				}
				return nil
			},

			func(ctx context.Context, key string) error {
				contains := strings.Contains(key, GranteeIDPrefix)
				if contains {
					if secondLoadWithPrefixReturn.Load() {
						return nil
					}
					return errors.New("mock loadwithprefix error")
				}

				if firstLoadWithPrefixReturn.Load() {
					return nil
				}
				return errors.New("mock loadWithPrefix error")
			},
		)

		tests := []struct {
			isValid      bool
			firstReturn  bool
			secondReturn bool

			description string
		}{
			{true, true, true, "valid both load with prefix success"},
			{false, true, false, "invalid first loadWithPrefix success second fail"},
			{false, false, true, "invalid first loadWithPrefix fail and second success"},
			{false, false, false, "invalid first loadWithPrefix fail and second fail"},
		}

		for _, test := range tests {
			t.Run(test.description, func(t *testing.T) {
				firstLoadWithPrefixReturn.Store(test.firstReturn)
				secondLoadWithPrefixReturn.Store(test.secondReturn)

				policy, err := c.ListPolicy(ctx, tenant)
				if test.isValid {
					assert.NoError(t, err)
					assert.Equal(t, 6, len(policy))
					ps := []*milvuspb.GrantEntity{
						grant("role1", "obj1", "obj_name1", "PrivilegeLoad", "default"),
						grant("role1", "obj1", "obj_name1", "PrivilegeRelease", "default"),
						grant("role1", "obj1", "obj_name1", util.AnyWord, "default"),
						grant("role2", "obj2", "obj_name2", "PrivilegeLoad", "default"),
						grant("role2", "obj2", "obj_name2", "PrivilegeRelease", "default"),
						grant("role2", "obj2", "obj_name2", util.AnyWord, "default"),
					}
					assert.ElementsMatch(t, ps, policy)
				} else {
					assert.Error(t, err)
				}
			})
		}
	})
}

func TestRBAC_Backup(t *testing.T) {
	etcdCli, _ := etcd.GetEtcdClient(
		Params.EtcdCfg.UseEmbedEtcd.GetAsBool(),
		Params.EtcdCfg.EtcdUseSSL.GetAsBool(),
		Params.EtcdCfg.Endpoints.GetAsStrings(),
		Params.EtcdCfg.EtcdTLSCert.GetValue(),
		Params.EtcdCfg.EtcdTLSKey.GetValue(),
		Params.EtcdCfg.EtcdTLSCACert.GetValue(),
		Params.EtcdCfg.EtcdTLSMinVersion.GetValue())
	rootPath := "/test/rbac"
	metaKV := etcdkv.NewEtcdKV(etcdCli, rootPath)
	defer metaKV.RemoveWithPrefix(context.TODO(), "")
	defer metaKV.Close()
	c := NewCatalog(metaKV, nil)

	ctx := context.Background()
	c.CreateRole(ctx, util.DefaultTenant, &milvuspb.RoleEntity{Name: "role1"})
	c.AlterGrant(ctx, util.DefaultTenant, &milvuspb.GrantEntity{
		Role:       &milvuspb.RoleEntity{Name: "role1"},
		Object:     &milvuspb.ObjectEntity{Name: "obj1"},
		ObjectName: "obj_name1",
		DbName:     util.DefaultDBName,
		Grantor: &milvuspb.GrantorEntity{
			User:      &milvuspb.UserEntity{Name: "user1"},
			Privilege: &milvuspb.PrivilegeEntity{Name: "PrivilegeLoad"},
		},
	}, milvuspb.OperatePrivilegeType_Grant)
	c.AlterCredential(ctx, &model.Credential{
		Username:          "user1",
		EncryptedPassword: "passwd",
	})
	c.AlterUserRole(ctx, util.DefaultTenant, &milvuspb.UserEntity{Name: "user1"}, &milvuspb.RoleEntity{Name: "role1"}, milvuspb.OperateUserRoleType_AddUserToRole)

	c.SavePrivilegeGroup(ctx, &milvuspb.PrivilegeGroupInfo{
		GroupName:  "custom_group",
		Privileges: []*milvuspb.PrivilegeEntity{{Name: "CreateCollection"}},
	})

	// test backup success
	backup, err := c.BackupRBAC(ctx, util.DefaultTenant)
	assert.NoError(t, err)
	assert.Equal(t, 1, len(backup.Grants))
	assert.Equal(t, "obj_name1", backup.Grants[0].ObjectName)
	assert.Equal(t, "role1", backup.Grants[0].Role.Name)
	assert.Equal(t, 1, len(backup.Users))
	assert.Equal(t, "user1", backup.Users[0].User)
	assert.Equal(t, 1, len(backup.Users[0].Roles))
	assert.Equal(t, 1, len(backup.Roles))
	assert.Equal(t, 1, len(backup.PrivilegeGroups))
	assert.Equal(t, "custom_group", backup.PrivilegeGroups[0].GroupName)
	assert.Equal(t, "CreateCollection", backup.PrivilegeGroups[0].Privileges[0].Name)
}

func TestRBAC_Restore(t *testing.T) {
	etcdCli, _ := etcd.GetEtcdClient(
		Params.EtcdCfg.UseEmbedEtcd.GetAsBool(),
		Params.EtcdCfg.EtcdUseSSL.GetAsBool(),
		Params.EtcdCfg.Endpoints.GetAsStrings(),
		Params.EtcdCfg.EtcdTLSCert.GetValue(),
		Params.EtcdCfg.EtcdTLSKey.GetValue(),
		Params.EtcdCfg.EtcdTLSCACert.GetValue(),
		Params.EtcdCfg.EtcdTLSMinVersion.GetValue())
	rootPath := "/test/rbac"
	metaKV := etcdkv.NewEtcdKV(etcdCli, rootPath)
	defer metaKV.RemoveWithPrefix(context.TODO(), "")
	defer metaKV.Close()
	c := NewCatalog(metaKV, nil)

	ctx := context.Background()

	rbacMeta := &milvuspb.RBACMeta{
		Users: []*milvuspb.UserInfo{
			{
				User:     "user1",
				Password: "passwd",
				Roles: []*milvuspb.RoleEntity{
					{
						Name: "role1",
					},
				},
			},
		},
		Roles: []*milvuspb.RoleEntity{
			{
				Name: "role1",
			},
		},

		Grants: []*milvuspb.GrantEntity{
			{
				Role:       &milvuspb.RoleEntity{Name: "role1"},
				Object:     &milvuspb.ObjectEntity{Name: "obj1"},
				ObjectName: "obj_name1",
				DbName:     util.DefaultDBName,
				Grantor: &milvuspb.GrantorEntity{
					User:      &milvuspb.UserEntity{Name: "user1"},
					Privilege: &milvuspb.PrivilegeEntity{Name: "Load"},
				},
			},
		},

		PrivilegeGroups: []*milvuspb.PrivilegeGroupInfo{
			{
				GroupName:  "custom_group",
				Privileges: []*milvuspb.PrivilegeEntity{{Name: "CreateCollection"}},
			},
		},
	}
	// test restore success
	err := c.RestoreRBAC(ctx, util.DefaultTenant, rbacMeta)
	assert.NoError(t, err)

	// check user
	users, err := c.ListCredentialsWithPasswd(ctx)
	assert.NoError(t, err)
	assert.Len(t, users, 1)
	assert.Equal(t, users["user1"], "passwd")
	// check role
	roles, err := c.ListRole(ctx, util.DefaultTenant, nil, false)
	assert.NoError(t, err)
	assert.Len(t, roles, 1)
	assert.Equal(t, "role1", roles[0].Role.Name)
	// check grant
	grants, err := c.ListGrant(ctx, util.DefaultTenant, &milvuspb.GrantEntity{
		Role:   roles[0].Role,
		DbName: util.AnyWord,
	})
	assert.NoError(t, err)
	assert.Len(t, grants, 1)
	assert.Equal(t, "obj_name1", grants[0].ObjectName)
	assert.Equal(t, "role1", grants[0].Role.Name)
	assert.Equal(t, "user1", grants[0].Grantor.User.Name)
	assert.Equal(t, "Load", grants[0].Grantor.Privilege.Name)
	// check privilege group
	privGroups, err := c.ListPrivilegeGroups(ctx)
	assert.NoError(t, err)
	assert.Len(t, privGroups, 1)
	assert.Equal(t, "custom_group", privGroups[0].GroupName)
	assert.Equal(t, "CreateCollection", privGroups[0].Privileges[0].Name)

	rbacMeta2 := &milvuspb.RBACMeta{
		Users: []*milvuspb.UserInfo{
			{
				User:     "user2",
				Password: "passwd",
				Roles: []*milvuspb.RoleEntity{
					{
						Name: "role2",
					},
				},
			},
		},
		Roles: []*milvuspb.RoleEntity{
			{
				Name: "role2",
			},
		},

		Grants: []*milvuspb.GrantEntity{
			{
				Role:       &milvuspb.RoleEntity{Name: "role2"},
				Object:     &milvuspb.ObjectEntity{Name: "obj2"},
				ObjectName: "obj_name2",
				DbName:     util.DefaultDBName,
				Grantor: &milvuspb.GrantorEntity{
					User:      &milvuspb.UserEntity{Name: "user2"},
					Privilege: &milvuspb.PrivilegeEntity{Name: "Load"},
				},
			},
		},

		PrivilegeGroups: []*milvuspb.PrivilegeGroupInfo{
			{
				GroupName:  "custom_group2",
				Privileges: []*milvuspb.PrivilegeEntity{{Name: "DropCollection"}},
			},
		},
	}

	// test restore failed and roll back
	err = c.RestoreRBAC(ctx, util.DefaultTenant, rbacMeta2)
	assert.NoError(t, err)

	// check user
	users, err = c.ListCredentialsWithPasswd(ctx)
	assert.NoError(t, err)
	assert.Len(t, users, 2)
	// check role
	roles, err = c.ListRole(ctx, util.DefaultTenant, nil, false)
	assert.NoError(t, err)
	assert.Len(t, roles, 2)
	// check grant
	grants, err = c.ListGrant(ctx, util.DefaultTenant, &milvuspb.GrantEntity{
		Role:   &milvuspb.RoleEntity{Name: "role2"},
		DbName: util.AnyWord,
	})
	assert.NoError(t, err)
	assert.Len(t, grants, 1)
	assert.Equal(t, "obj_name2", grants[0].ObjectName)
	assert.Equal(t, "role2", grants[0].Role.Name)
	assert.Equal(t, "user2", grants[0].Grantor.User.Name)
	assert.Equal(t, "Load", grants[0].Grantor.Privilege.Name)
	// check privilege group
	privGroups, err = c.ListPrivilegeGroups(ctx)
	assert.NoError(t, err)
	assert.Len(t, privGroups, 2)
}

func TestRBAC_PrivilegeGroup(t *testing.T) {
	ctx := context.TODO()
	group1 := "group1"
	group2 := "group2"
	key1 := BuildPrivilegeGroupkey(group1)
	key2 := BuildPrivilegeGroupkey(group2)
	privGroupInfo1 := &milvuspb.PrivilegeGroupInfo{GroupName: group1, Privileges: []*milvuspb.PrivilegeEntity{{Name: "priv10"}, {Name: "priv11"}}}
	privGroupInfo2 := &milvuspb.PrivilegeGroupInfo{GroupName: group2, Privileges: []*milvuspb.PrivilegeEntity{{Name: "priv20"}, {Name: "priv21"}}}
	v1, _ := proto.Marshal(privGroupInfo1)
	v2, _ := proto.Marshal(privGroupInfo2)

	t.Run("test GetPrivilegeGroup", func(t *testing.T) {
		var (
			kvmock = mocks.NewTxnKV(t)
			c      = NewCatalog(kvmock, nil)
		)
		kvmock.EXPECT().Load(mock.Anything, key1).Return(string(v1), nil)
		kvmock.EXPECT().Load(mock.Anything, key2).Return("", merr.ErrIoKeyNotFound)

		tests := []struct {
			description        string
			expectedErr        error
			groupName          string
			expectedPrivileges []string
		}{
			{"group not found", fmt.Errorf("privilege group [%s] does not exist", group2), group2, nil},
			{"valid group", nil, group1, []string{"priv10", "priv11"}},
		}
		for _, test := range tests {
			t.Run(test.description, func(t *testing.T) {
				group, err := c.GetPrivilegeGroup(ctx, test.groupName)
				if test.expectedErr != nil {
					assert.Error(t, err, test.expectedErr)
				} else {
					assert.NoError(t, err)
					assert.ElementsMatch(t, getPrivilegeNames(group.Privileges), test.expectedPrivileges)
				}
			})
		}
	})

	t.Run("test DropPrivilegeGroup", func(t *testing.T) {
		var (
			kvmock = mocks.NewTxnKV(t)
			c      = NewCatalog(kvmock, nil)
		)

		kvmock.EXPECT().Remove(mock.Anything, key1).Return(nil)
		kvmock.EXPECT().Remove(mock.Anything, key2).Return(errors.New("Mock remove failure"))

		tests := []struct {
			description string
			isValid     bool
			groupName   string
		}{
			{"valid group", true, group1},
			{"remove failure", false, group2},
		}

		for _, test := range tests {
			t.Run(test.description, func(t *testing.T) {
				err := c.DropPrivilegeGroup(ctx, test.groupName)
				if test.isValid {
					assert.NoError(t, err)
				} else {
					assert.Error(t, err)
				}
			})
		}
	})

	t.Run("test SavePrivilegeGroup", func(t *testing.T) {
		var (
			kvmock = mocks.NewTxnKV(t)
			c      = NewCatalog(kvmock, nil)
		)

		kvmock.EXPECT().Save(mock.Anything, key1, mock.Anything).Return(nil)
		kvmock.EXPECT().Save(mock.Anything, key2, mock.Anything).Return(nil)

		tests := []struct {
			description string
			isValid     bool
			group       *milvuspb.PrivilegeGroupInfo
		}{
			{"valid group with existing key", true, &milvuspb.PrivilegeGroupInfo{GroupName: group1, Privileges: []*milvuspb.PrivilegeEntity{{Name: "priv10"}, {Name: "priv11"}}}},
			{"valid group without existing key", true, &milvuspb.PrivilegeGroupInfo{GroupName: group2, Privileges: []*milvuspb.PrivilegeEntity{{Name: "priv10"}, {Name: "priv11"}}}},
		}

		for _, test := range tests {
			t.Run(test.description, func(t *testing.T) {
				err := c.SavePrivilegeGroup(ctx, test.group)
				if test.isValid {
					assert.NoError(t, err)
				} else {
					assert.Error(t, err)
				}
			})
		}
	})

	t.Run("test ListPrivilegeGroups", func(t *testing.T) {
		var (
			kvmock = mocks.NewTxnKV(t)
			c      = NewCatalog(kvmock, nil)
		)

		kvmock.EXPECT().LoadWithPrefix(mock.Anything, PrivilegeGroupPrefix).Return(
			[]string{key1, key2},
			[]string{string(v1), string(v2)},
			nil,
		)
		groups, err := c.ListPrivilegeGroups(ctx)
		assert.NoError(t, err)
		groupNames := lo.Map(groups, func(g *milvuspb.PrivilegeGroupInfo, _ int) string {
			return g.GroupName
		})
		assert.ElementsMatch(t, groupNames, []string{group1, group2})
		assert.ElementsMatch(t, getPrivilegeNames(groups[0].Privileges), []string{"priv10", "priv11"})
		assert.ElementsMatch(t, getPrivilegeNames(groups[1].Privileges), []string{"priv20", "priv21"})
	})
}

func getPrivilegeNames(privileges []*milvuspb.PrivilegeEntity) []string {
	if len(privileges) == 0 {
		return []string{}
	}
	return lo.Map(privileges, func(p *milvuspb.PrivilegeEntity, _ int) string {
		return p.Name
	})
}

func TestCatalog_AlterDatabase(t *testing.T) {
	kvmock := mocks.NewSnapShotKV(t)
	c := NewCatalog(nil, kvmock)
	db := model.NewDatabase(1, "db", pb.DatabaseState_DatabaseCreated, nil)

	kvmock.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
	ctx := context.Background()

	// test alter database success
	newDB := db.Clone()
	db.Properties = []*commonpb.KeyValuePair{
		{
			Key:   "key1",
			Value: "value1",
		},
	}
	err := c.AlterDatabase(ctx, newDB, typeutil.ZeroTimestamp)
	assert.NoError(t, err)

	// test alter database fail
	mockErr := errors.New("access kv store error")
	kvmock.ExpectedCalls = nil
	kvmock.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(mockErr)
	err = c.AlterDatabase(ctx, newDB, typeutil.ZeroTimestamp)
	assert.ErrorIs(t, err, mockErr)
}

func TestCatalog_listFunctionError(t *testing.T) {
	mockSnapshot := newMockSnapshot(t)
	kc := NewCatalog(nil, mockSnapshot).(*Catalog)
	mockSnapshot.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return(nil, nil, errors.New("mock error"))
	_, err := kc.listFunctions(context.TODO(), 1, 1)
	assert.Error(t, err)

	mockSnapshot.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything, mock.Anything).Return([]string{"test-key"}, []string{"invalid bytes"}, nil)
	_, err = kc.listFunctions(context.TODO(), 1, 1)
	assert.Error(t, err)
}

// TestRBACPrefixMatch tests the fix for RBAC prefix matching issues.
// The fix ensures that when querying by prefix (e.g., user1), it won't
// mistakenly match other entries with similar prefixes (e.g., user10).
func TestRBACPrefixMatch(t *testing.T) {
	ctx := context.Background()
	tenant := util.DefaultTenant

	t.Run("getRolesByUsername should not match similar prefix usernames", func(t *testing.T) {
		// Scenario: user1 and user10 exist, querying user1 should not return user10's roles
		kvmock := mocks.NewTxnKV(t)
		c := NewCatalog(kvmock, nil).(*Catalog)

		// The fix adds "/" to the prefix, so query for "user1" uses prefix "RoleMappingPrefix/tenant/user1/"
		// This ensures user10's data won't be matched
		user1Prefix := funcutil.HandleTenantForEtcdKey(RoleMappingPrefix, tenant, "user1") + "/"

		// Mock: return keys that match the prefix correctly
		// Key format: prefix + roleName (where prefix already ends with "/")
		kvmock.EXPECT().LoadWithPrefix(mock.Anything, user1Prefix).Return(
			[]string{
				user1Prefix + "roleA",
				user1Prefix + "roleB",
			},
			nil,
			nil,
		)

		roles, err := c.getRolesByUsername(ctx, tenant, "user1")
		assert.NoError(t, err)
		assert.Equal(t, 2, len(roles))
		assert.Contains(t, roles, "roleA")
		assert.Contains(t, roles, "roleB")
	})

	t.Run("ListGrant should not match similar prefix role names", func(t *testing.T) {
		// Scenario: role1 and role10 exist, querying role1 should not return role10's grants
		kvmock := mocks.NewTxnKV(t)
		c := NewCatalog(kvmock, nil).(*Catalog)

		// The fix adds "/" to the prefix for role query
		role1Prefix := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant, "role1") + "/"

		// Mock: return only role1's grants (the fix ensures role10 won't be matched)
		// Key format: prefix + object + "/" + objectName (prefix already ends with "/")
		kvmock.EXPECT().LoadWithPrefix(mock.Anything, role1Prefix).Return(
			[]string{
				role1Prefix + "Collection/col1",
				role1Prefix + "Collection/col2",
			},
			[]string{
				"grantee_id_1",
				"grantee_id_2",
			},
			nil,
		)

		// Mock granteeID lookups with "/" suffix
		granteeID1Prefix := funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, "grantee_id_1") + "/"
		granteeID2Prefix := funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, "grantee_id_2") + "/"

		kvmock.EXPECT().LoadWithPrefix(mock.Anything, granteeID1Prefix).Return(
			[]string{granteeID1Prefix + "Load"},
			[]string{"admin"},
			nil,
		)
		kvmock.EXPECT().LoadWithPrefix(mock.Anything, granteeID2Prefix).Return(
			[]string{granteeID2Prefix + "Release"},
			[]string{"admin"},
			nil,
		)

		entity := &milvuspb.GrantEntity{
			Role:   &milvuspb.RoleEntity{Name: "role1"},
			DbName: util.DefaultDBName,
		}
		grants, err := c.ListGrant(ctx, tenant, entity)
		assert.NoError(t, err)
		assert.Equal(t, 2, len(grants))
		for _, g := range grants {
			assert.Equal(t, "role1", g.GetRole().GetName())
		}
	})

	t.Run("ListGrant granteeID prefix should not match similar prefixes", func(t *testing.T) {
		// Scenario: granteeID "abc" and "abc123" exist, querying abc should not return abc123's privileges
		kvmock := mocks.NewTxnKV(t)
		c := NewCatalog(kvmock, nil).(*Catalog)

		// Setup for a specific object query that triggers appendGrantEntity
		granteeKey := funcutil.HandleTenantForEtcdKey(GranteePrefix, tenant,
			fmt.Sprintf("%s/%s/%s", "testRole", "Collection", "testCol"))

		// Mock Load to return a granteeID
		kvmock.EXPECT().Load(mock.Anything, granteeKey).Return("grantee_abc", nil)

		// The fix adds "/" to granteeID prefix query
		granteeIDPrefix := funcutil.HandleTenantForEtcdKey(GranteeIDPrefix, tenant, "grantee_abc") + "/"

		// Mock: return only grantee_abc's privileges (not grantee_abc123's)
		kvmock.EXPECT().LoadWithPrefix(mock.Anything, granteeIDPrefix).Return(
			[]string{
				granteeIDPrefix + "Insert",
				granteeIDPrefix + "Delete",
			},
			[]string{"user1", "user2"},
			nil,
		)

		entity := &milvuspb.GrantEntity{
			Role:       &milvuspb.RoleEntity{Name: "testRole"},
			Object:     &milvuspb.ObjectEntity{Name: "Collection"},
			ObjectName: "testCol",
			DbName:     util.DefaultDBName,
		}
		grants, err := c.ListGrant(ctx, tenant, entity)
		assert.NoError(t, err)
		assert.Equal(t, 2, len(grants))
		// Verify privileges are correctly parsed
		privileges := make([]string, 0, len(grants))
		for _, g := range grants {
			privileges = append(privileges, g.GetGrantor().GetPrivilege().GetName())
		}
		assert.Contains(t, privileges, util.PrivilegeNameForAPI("Insert"))
		assert.Contains(t, privileges, util.PrivilegeNameForAPI("Delete"))
	})
}

func TestCatalog_FileResource(t *testing.T) {
	ctx := context.Background()
	mockErr := errors.New("mock error")

	t.Run("BuildFileResourceKey", func(t *testing.T) {
		key := BuildFileResourceKey(12345)
		assert.Equal(t, fmt.Sprintf("%s/%d", FileResourceMetaPrefix, 12345), key)
	})

	t.Run("SaveFileResource success", func(t *testing.T) {
		kvmock := mocks.NewTxnKV(t)
		c := NewCatalog(kvmock, nil)

		resource := &internalpb.FileResourceInfo{
			Id:   1,
			Name: "test_resource",
			Path: "/path/to/resource",
		}

		kvmock.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(nil)

		err := c.SaveFileResource(ctx, resource, 1)
		assert.NoError(t, err)
	})

	t.Run("SaveFileResource write error", func(t *testing.T) {
		kvmock := mocks.NewTxnKV(t)
		c := NewCatalog(kvmock, nil)

		// nil resource will cause proto marshal to handle gracefully, test MultiSave error instead
		resource := &internalpb.FileResourceInfo{
			Id:   1,
			Name: "test_resource",
			Path: "/path/to/resource",
		}

		kvmock.EXPECT().MultiSave(mock.Anything, mock.Anything).Return(mockErr)

		err := c.SaveFileResource(ctx, resource, 1)
		assert.Error(t, err)
	})

	t.Run("RemoveFileResource success", func(t *testing.T) {
		kvmock := mocks.NewTxnKV(t)
		c := NewCatalog(kvmock, nil)

		kvmock.EXPECT().MultiSaveAndRemove(mock.Anything, mock.Anything, mock.Anything).Return(nil)

		err := c.RemoveFileResource(ctx, 1, 1)
		assert.NoError(t, err)
	})

	t.Run("RemoveFileResource error", func(t *testing.T) {
		kvmock := mocks.NewTxnKV(t)
		c := NewCatalog(kvmock, nil)

		kvmock.EXPECT().MultiSaveAndRemove(mock.Anything, mock.Anything, mock.Anything).Return(mockErr)

		err := c.RemoveFileResource(ctx, 1, 1)
		assert.Error(t, err)
	})

	t.Run("ListFileResource success with version", func(t *testing.T) {
		kvmock := mocks.NewTxnKV(t)
		c := NewCatalog(kvmock, nil)

		resource := &internalpb.FileResourceInfo{
			Id:   1,
			Name: "test_resource",
			Path: "/path/to/resource",
		}
		value, err := proto.Marshal(resource)
		assert.NoError(t, err)

		kvmock.EXPECT().LoadWithPrefix(mock.Anything, FileResourceMetaPrefix).Return(
			[]string{BuildFileResourceKey(1)},
			[]string{string(value)},
			nil,
		)
		kvmock.EXPECT().Has(mock.Anything, FileResourceVersionKey).Return(true, nil)
		kvmock.EXPECT().Load(mock.Anything, FileResourceVersionKey).Return("123", nil)

		resources, version, err := c.ListFileResource(ctx)
		assert.NoError(t, err)
		assert.Equal(t, uint64(123), version)
		assert.Equal(t, 1, len(resources))
		assert.Equal(t, int64(1), resources[0].Id)
		assert.Equal(t, "test_resource", resources[0].Name)
		assert.Equal(t, "/path/to/resource", resources[0].Path)
	})

	t.Run("ListFileResource success without version", func(t *testing.T) {
		kvmock := mocks.NewTxnKV(t)
		c := NewCatalog(kvmock, nil)

		resource := &internalpb.FileResourceInfo{
			Id:   2,
			Name: "test_resource_2",
			Path: "/path/to/resource2",
		}
		value, err := proto.Marshal(resource)
		assert.NoError(t, err)

		kvmock.EXPECT().LoadWithPrefix(mock.Anything, FileResourceMetaPrefix).Return(
			[]string{BuildFileResourceKey(2)},
			[]string{string(value)},
			nil,
		)
		kvmock.EXPECT().Has(mock.Anything, FileResourceVersionKey).Return(false, nil)

		resources, version, err := c.ListFileResource(ctx)
		assert.NoError(t, err)
		assert.Equal(t, uint64(0), version)
		assert.Equal(t, 1, len(resources))
	})

	t.Run("ListFileResource LoadWithPrefix error", func(t *testing.T) {
		kvmock := mocks.NewTxnKV(t)
		c := NewCatalog(kvmock, nil)

		kvmock.EXPECT().LoadWithPrefix(mock.Anything, FileResourceMetaPrefix).Return(nil, nil, mockErr)

		resources, version, err := c.ListFileResource(ctx)
		assert.Error(t, err)
		assert.Nil(t, resources)
		assert.Equal(t, uint64(0), version)
	})

	t.Run("ListFileResource Has error", func(t *testing.T) {
		kvmock := mocks.NewTxnKV(t)
		c := NewCatalog(kvmock, nil)

		kvmock.EXPECT().LoadWithPrefix(mock.Anything, FileResourceMetaPrefix).Return([]string{}, []string{}, nil)
		kvmock.EXPECT().Has(mock.Anything, FileResourceVersionKey).Return(false, mockErr)

		resources, version, err := c.ListFileResource(ctx)
		assert.Error(t, err)
		assert.Nil(t, resources)
		assert.Equal(t, uint64(0), version)
	})

	t.Run("ListFileResource Load version error", func(t *testing.T) {
		kvmock := mocks.NewTxnKV(t)
		c := NewCatalog(kvmock, nil)

		kvmock.EXPECT().LoadWithPrefix(mock.Anything, FileResourceMetaPrefix).Return([]string{}, []string{}, nil)
		kvmock.EXPECT().Has(mock.Anything, FileResourceVersionKey).Return(true, nil)
		kvmock.EXPECT().Load(mock.Anything, FileResourceVersionKey).Return("", mockErr)

		resources, version, err := c.ListFileResource(ctx)
		assert.Error(t, err)
		assert.Nil(t, resources)
		assert.Equal(t, uint64(0), version)
	})

	t.Run("ListFileResource parse version error", func(t *testing.T) {
		kvmock := mocks.NewTxnKV(t)
		c := NewCatalog(kvmock, nil)

		kvmock.EXPECT().LoadWithPrefix(mock.Anything, FileResourceMetaPrefix).Return([]string{}, []string{}, nil)
		kvmock.EXPECT().Has(mock.Anything, FileResourceVersionKey).Return(true, nil)
		kvmock.EXPECT().Load(mock.Anything, FileResourceVersionKey).Return("invalid_version", nil)

		resources, version, err := c.ListFileResource(ctx)
		assert.Error(t, err)
		assert.Nil(t, resources)
		assert.Equal(t, uint64(0), version)
	})

	t.Run("ListFileResource unmarshal error", func(t *testing.T) {
		kvmock := mocks.NewTxnKV(t)
		c := NewCatalog(kvmock, nil)

		kvmock.EXPECT().LoadWithPrefix(mock.Anything, FileResourceMetaPrefix).Return(
			[]string{BuildFileResourceKey(1)},
			[]string{"invalid_proto_data"},
			nil,
		)
		kvmock.EXPECT().Has(mock.Anything, FileResourceVersionKey).Return(false, nil)

		resources, version, err := c.ListFileResource(ctx)
		assert.Error(t, err)
		assert.Nil(t, resources)
		assert.Equal(t, uint64(0), version)
	})

	t.Run("ListFileResource empty", func(t *testing.T) {
		kvmock := mocks.NewTxnKV(t)
		c := NewCatalog(kvmock, nil)

		kvmock.EXPECT().LoadWithPrefix(mock.Anything, FileResourceMetaPrefix).Return([]string{}, []string{}, nil)
		kvmock.EXPECT().Has(mock.Anything, FileResourceVersionKey).Return(false, nil)

		resources, version, err := c.ListFileResource(ctx)
		assert.NoError(t, err)
		assert.Equal(t, uint64(0), version)
		assert.Equal(t, 0, len(resources))
	})

	t.Run("ListFileResource multiple resources", func(t *testing.T) {
		kvmock := mocks.NewTxnKV(t)
		c := NewCatalog(kvmock, nil)

		resource1 := &internalpb.FileResourceInfo{
			Id:   1,
			Name: "resource1",
			Path: "/path/1",
		}
		resource2 := &internalpb.FileResourceInfo{
			Id:   2,
			Name: "resource2",
			Path: "/path/2",
		}
		value1, _ := proto.Marshal(resource1)
		value2, _ := proto.Marshal(resource2)

		kvmock.EXPECT().LoadWithPrefix(mock.Anything, FileResourceMetaPrefix).Return(
			[]string{BuildFileResourceKey(1), BuildFileResourceKey(2)},
			[]string{string(value1), string(value2)},
			nil,
		)
		kvmock.EXPECT().Has(mock.Anything, FileResourceVersionKey).Return(true, nil)
		kvmock.EXPECT().Load(mock.Anything, FileResourceVersionKey).Return("456", nil)

		resources, version, err := c.ListFileResource(ctx)
		assert.NoError(t, err)
		assert.Equal(t, uint64(456), version)
		assert.Equal(t, 2, len(resources))
	})
}
