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

516 lines
14 KiB
Go
Raw Blame History

// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package model
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/pkg/v2/common"
pb "github.com/milvus-io/milvus/pkg/v2/proto/etcdpb"
)
var (
colID int64 = 1
colName = "c"
fieldID int64 = 101
fieldName = "field110"
partID int64 = 20
partName = "testPart"
tenantID = "tenant-1"
functionID int64 = 1
functionName = "test-bm25"
typeParams = []*commonpb.KeyValuePair{
{
Key: "field110-k1",
Value: "field110-v1",
},
}
startPositions = []*commonpb.KeyDataPair{
{
Key: "k1",
Data: []byte{byte(1)},
},
}
colModel = &Collection{
TenantID: tenantID,
CollectionID: colID,
Name: colName,
AutoID: false,
Description: "none",
Fields: []*Field{fieldModel},
VirtualChannelNames: []string{"vch"},
PhysicalChannelNames: []string{"pch"},
ShardsNum: common.DefaultShardsNum,
CreateTime: 1,
StartPositions: startPositions,
ConsistencyLevel: commonpb.ConsistencyLevel_Strong,
Partitions: []*Partition{
{
PartitionID: partID,
PartitionName: partName,
PartitionCreatedTimestamp: 1,
},
},
Properties: []*commonpb.KeyValuePair{
{
Key: "k",
Value: "v",
},
},
}
deprecatedColPb = &pb.CollectionInfo{
ID: colID,
Schema: &schemapb.CollectionSchema{
Name: colName,
Description: "none",
AutoID: false,
Fields: []*schemapb.FieldSchema{filedSchemaPb},
},
CreateTime: 1,
PartitionIDs: []int64{partID},
PartitionNames: []string{partName},
PartitionCreatedTimestamps: []uint64{1},
FieldIndexes: []*pb.FieldIndexInfo{
{
FiledID: fieldID,
IndexID: indexID,
},
},
VirtualChannelNames: []string{"vch"},
PhysicalChannelNames: []string{"pch"},
ShardsNum: common.DefaultShardsNum,
StartPositions: startPositions,
ConsistencyLevel: commonpb.ConsistencyLevel_Strong,
Properties: []*commonpb.KeyValuePair{
{
Key: "k",
Value: "v",
},
},
}
)
func TestUnmarshalCollectionModel(t *testing.T) {
ret := UnmarshalCollectionModel(deprecatedColPb)
ret.TenantID = tenantID
assert.Equal(t, ret, colModel)
assert.Nil(t, UnmarshalCollectionModel(nil))
}
func TestMarshalCollectionModel(t *testing.T) {
assert.Nil(t, MarshalCollectionModel(nil))
}
func TestCollection_GetPartitionNum(t *testing.T) {
coll := &Collection{
Partitions: []*Partition{
{State: pb.PartitionState_PartitionCreated},
{State: pb.PartitionState_PartitionCreating},
{State: pb.PartitionState_PartitionCreated},
{State: pb.PartitionState_PartitionDropping},
{State: pb.PartitionState_PartitionCreated},
{State: pb.PartitionState_PartitionDropped},
},
}
assert.Equal(t, 3, coll.GetPartitionNum(true))
assert.Equal(t, 6, coll.GetPartitionNum(false))
}
func TestCollection_Equal(t *testing.T) {
equal := func(a, b Collection) bool {
return a.Equal(b)
}
type args struct {
a, b Collection
}
tests := []struct {
name string
args args
want bool
}{
{
args: args{
a: Collection{TenantID: "aaa"},
b: Collection{TenantID: "bbb"},
},
want: false,
},
{
args: args{
a: Collection{
TenantID: "aaa",
Partitions: []*Partition{{PartitionName: "default"}},
},
b: Collection{
TenantID: "aaa",
},
},
want: false,
},
{
args: args{
a: Collection{
TenantID: "aaa",
Partitions: []*Partition{{PartitionName: "default"}},
Name: "aaa",
},
b: Collection{
TenantID: "aaa",
Partitions: []*Partition{{PartitionName: "default"}},
Name: "bbb",
},
},
want: false,
},
{
args: args{
a: Collection{
TenantID: "aaa",
Partitions: []*Partition{{PartitionName: "default"}},
Name: "aaa",
Description: "aaa",
},
b: Collection{
TenantID: "aaa",
Partitions: []*Partition{{PartitionName: "default"}},
Name: "aaa",
Description: "bbb",
},
},
want: false,
},
{
args: args{
a: Collection{
TenantID: "aaa",
Partitions: []*Partition{{PartitionName: "default"}},
Name: "aaa",
Description: "aaa",
AutoID: false,
},
b: Collection{
TenantID: "aaa",
Partitions: []*Partition{{PartitionName: "default"}},
Name: "aaa",
Description: "aaa",
AutoID: true,
},
},
want: false,
},
{
args: args{
a: Collection{
TenantID: "aaa",
Partitions: []*Partition{{PartitionName: "default"}},
Name: "aaa",
Description: "aaa",
AutoID: false,
Fields: []*Field{{Name: "f1", TypeParams: []*commonpb.KeyValuePair{
{Key: "dim", Value: "128"},
}}},
},
b: Collection{
TenantID: "aaa",
Partitions: []*Partition{{PartitionName: "default"}},
Name: "aaa",
Description: "aaa",
AutoID: false,
Fields: []*Field{{Name: "f1", TypeParams: []*commonpb.KeyValuePair{
{Key: "dim", Value: "256"},
}}},
},
},
want: false,
},
{
args: args{
a: Collection{
TenantID: "aaa",
Partitions: []*Partition{{PartitionName: "default"}},
Name: "aaa",
Description: "aaa",
AutoID: false,
Fields: []*Field{{Name: "f1", TypeParams: []*commonpb.KeyValuePair{
{Key: "dim", Value: "128"},
}}},
ShardsNum: 1,
},
b: Collection{
TenantID: "aaa",
Partitions: []*Partition{{PartitionName: "default"}},
Name: "aaa",
Description: "aaa",
AutoID: false,
Fields: []*Field{{Name: "f1", TypeParams: []*commonpb.KeyValuePair{
{Key: "dim", Value: "128"},
}}},
ShardsNum: 2,
},
},
want: false,
},
{
args: args{
a: Collection{
TenantID: "aaa",
Partitions: []*Partition{{PartitionName: "default"}},
Name: "aaa",
Description: "aaa",
AutoID: false,
Fields: []*Field{{Name: "f1", TypeParams: []*commonpb.KeyValuePair{
{Key: "dim", Value: "128"},
}}},
ShardsNum: 1,
ConsistencyLevel: commonpb.ConsistencyLevel_Strong,
},
b: Collection{
TenantID: "aaa",
Partitions: []*Partition{{PartitionName: "default"}},
Name: "aaa",
Description: "aaa",
AutoID: false,
Fields: []*Field{{Name: "f1", TypeParams: []*commonpb.KeyValuePair{
{Key: "dim", Value: "128"},
}}},
ShardsNum: 1,
ConsistencyLevel: commonpb.ConsistencyLevel_Bounded,
},
},
want: false,
},
{
args: args{
a: Collection{
TenantID: "aaa",
Partitions: []*Partition{{PartitionName: "default"}},
Name: "aaa",
Description: "aaa",
AutoID: false,
Fields: []*Field{{Name: "f1", TypeParams: []*commonpb.KeyValuePair{
{Key: "dim", Value: "128"},
}}},
ShardsNum: 1,
ConsistencyLevel: commonpb.ConsistencyLevel_Strong,
Properties: []*commonpb.KeyValuePair{
{Key: "ttl", Value: "200"},
},
},
b: Collection{
TenantID: "aaa",
Partitions: []*Partition{{PartitionName: "default"}},
Name: "aaa",
Description: "aaa",
AutoID: false,
Fields: []*Field{{Name: "f1", TypeParams: []*commonpb.KeyValuePair{
{Key: "dim", Value: "128"},
}}},
ShardsNum: 1,
ConsistencyLevel: commonpb.ConsistencyLevel_Strong,
Properties: []*commonpb.KeyValuePair{
{Key: "ttl", Value: "100"},
},
},
},
want: false,
},
{
args: args{
a: Collection{
TenantID: "aaa",
Partitions: []*Partition{{PartitionName: "default"}},
Name: "aaa",
Description: "aaa",
AutoID: false,
Fields: []*Field{{Name: "f1", TypeParams: []*commonpb.KeyValuePair{
{Key: "dim", Value: "128"},
}}},
ShardsNum: 1,
ConsistencyLevel: commonpb.ConsistencyLevel_Strong,
Properties: []*commonpb.KeyValuePair{
{Key: "ttl", Value: "200"},
},
EnableDynamicField: false,
},
b: Collection{
TenantID: "aaa",
Partitions: []*Partition{{PartitionName: "default"}},
Name: "aaa",
Description: "aaa",
AutoID: false,
Fields: []*Field{{Name: "f1", TypeParams: []*commonpb.KeyValuePair{
{Key: "dim", Value: "128"},
}}},
ShardsNum: 1,
ConsistencyLevel: commonpb.ConsistencyLevel_Strong,
Properties: []*commonpb.KeyValuePair{
{Key: "ttl", Value: "200"},
},
EnableDynamicField: true,
},
},
want: false,
},
{
args: args{
a: Collection{
TenantID: "aaa",
Partitions: []*Partition{{PartitionName: "default"}},
Name: "aaa",
Description: "aaa",
AutoID: false,
Fields: []*Field{{Name: "f1", TypeParams: []*commonpb.KeyValuePair{
{Key: "dim", Value: "128"},
}}},
ShardsNum: 1,
ConsistencyLevel: commonpb.ConsistencyLevel_Strong,
Properties: []*commonpb.KeyValuePair{
{Key: "ttl", Value: "200"},
},
EnableDynamicField: false,
},
b: Collection{
TenantID: "aaa",
Partitions: []*Partition{{PartitionName: "default"}},
Name: "aaa",
Description: "aaa",
AutoID: false,
Fields: []*Field{{Name: "f1", TypeParams: []*commonpb.KeyValuePair{
{Key: "dim", Value: "128"},
}}},
ShardsNum: 1,
ConsistencyLevel: commonpb.ConsistencyLevel_Strong,
Properties: []*commonpb.KeyValuePair{
{Key: "ttl", Value: "200"},
},
EnableDynamicField: false,
},
},
want: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := equal(tt.args.a, tt.args.b); got != tt.want {
t.Errorf("equal() = %v, want %v, collection a: %v, collection b: %v", got, tt.want, tt.args.a, tt.args.b)
}
})
}
}
func TestClone(t *testing.T) {
collection := &Collection{
TenantID: "1",
DBID: 2,
CollectionID: 3,
Partitions: []*Partition{
{
PartitionID: 4,
PartitionName: "5",
PartitionCreatedTimestamp: 6,
CollectionID: 7,
State: pb.PartitionState_PartitionCreated,
},
{
PartitionID: 8,
PartitionName: "9",
PartitionCreatedTimestamp: 10,
CollectionID: 11,
State: pb.PartitionState_PartitionCreating,
},
},
Name: "12",
DBName: "13",
Description: "14",
AutoID: true,
Fields: []*Field{
{
FieldID: 15,
Name: "16",
IsPrimaryKey: false,
Description: "17",
DataType: schemapb.DataType_Double,
TypeParams: []*commonpb.KeyValuePair{{Key: "18", Value: "19"}},
IndexParams: []*commonpb.KeyValuePair{{Key: "20", Value: "21"}},
AutoID: true,
State: schemapb.FieldState_FieldDropping,
IsDynamic: true,
IsPartitionKey: false,
IsClusteringKey: true,
IsFunctionOutput: false,
DefaultValue: nil,
ElementType: schemapb.DataType_String,
Nullable: true,
},
{
FieldID: 22,
Name: "23",
IsPrimaryKey: true,
Description: "24",
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{{Key: "25", Value: "26"}},
IndexParams: []*commonpb.KeyValuePair{{Key: "27", Value: "28"}},
AutoID: true,
State: schemapb.FieldState_FieldCreating,
IsDynamic: true,
IsPartitionKey: false,
IsClusteringKey: true,
IsFunctionOutput: false,
DefaultValue: nil,
ElementType: schemapb.DataType_VarChar,
Nullable: true,
},
},
Functions: []*Function{
{
ID: functionID,
Name: functionName,
Type: schemapb.FunctionType_BM25,
InputFieldIDs: []int64{101},
InputFieldNames: []string{"text"},
OutputFieldIDs: []int64{103},
OutputFieldNames: []string{"sparse"},
},
},
VirtualChannelNames: []string{"c1", "c2"},
PhysicalChannelNames: []string{"c3", "c4"},
ShardsNum: 2,
StartPositions: startPositions,
CreateTime: 1234,
ConsistencyLevel: commonpb.ConsistencyLevel_Eventually,
Aliases: []string{"a1", "a2"},
Properties: []*commonpb.KeyValuePair{{Key: "32", Value: "33"}},
State: pb.CollectionState_CollectionCreated,
EnableDynamicField: true,
}
clone1 := collection.Clone()
assert.Equal(t, clone1, collection)
clone2 := collection.ShallowClone()
assert.Equal(t, clone2, collection)
}
Reference in New Issue View Git Blame Copy Permalink