milvus/internal/indexnode/indexnode_test.go
wayblink 85e04d845c
Move GetComponentStates to milvus.proto, make it an external interface (#19594)
Signed-off-by: wayblink <anyang.wang@zilliz.com>

Signed-off-by: wayblink <anyang.wang@zilliz.com>
2022-10-10 15:55:22 +08:00

546 lines
15 KiB
Go

// 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 indexnode
import (
"context"
"os"
"testing"
"github.com/milvus-io/milvus/api/commonpb"
"github.com/stretchr/testify/assert"
)
//func TestRegister(t *testing.T) {
// var (
// factory = &mockFactory{}
// ctx = context.TODO()
// )
// Params.Init()
// in, err := NewIndexNode(ctx, factory)
// assert.Nil(t, err)
// in.SetEtcdClient(getEtcdClient())
// assert.Nil(t, in.initSession())
// assert.Nil(t, in.Register())
// key := in.session.ServerName
// if !in.session.Exclusive {
// key = fmt.Sprintf("%s-%d", key, in.session.ServerID)
// }
// resp, err := getEtcdClient().Get(ctx, path.Join(Params.EtcdCfg.MetaRootPath, sessionutil.DefaultServiceRoot, key))
// assert.Nil(t, err)
// assert.Equal(t, int64(1), resp.Count)
// sess := &sessionutil.Session{}
// assert.Nil(t, json.Unmarshal(resp.Kvs[0].Value, sess))
// assert.Equal(t, sess.ServerID, in.session.ServerID)
// assert.Equal(t, sess.Address, in.session.Address)
// assert.Equal(t, sess.ServerName, in.session.ServerName)
//
// // revoke lease
// in.session.Revoke(time.Second)
//
// in.chunkManager = storage.NewLocalChunkManager(storage.RootPath("/tmp/lib/milvus"))
// t.Run("CreateIndex FloatVector", func(t *testing.T) {
// var insertCodec storage.InsertCodec
//
// insertCodec.Schema = &etcdpb.CollectionMeta{
// ID: collectionID,
// Schema: &schemapb.CollectionSchema{
// Fields: []*schemapb.FieldSchema{
// {
// FieldID: floatVectorFieldID,
// Name: floatVectorFieldName,
// IsPrimaryKey: false,
// DataType: schemapb.DataType_FloatVector,
// },
// },
// },
// }
// data := make(map[UniqueID]storage.FieldData)
// tsData := make([]int64, nb)
// for i := 0; i < nb; i++ {
// tsData[i] = int64(i + 100)
// }
// data[tsFieldID] = &storage.Int64FieldData{
// NumRows: []int64{nb},
// Data: tsData,
// }
// data[floatVectorFieldID] = &storage.FloatVectorFieldData{
// NumRows: []int64{nb},
// Data: generateFloatVectors(),
// Dim: dim,
// }
// insertData := storage.InsertData{
// Data: data,
// Infos: []storage.BlobInfo{
// {
// Length: 10,
// },
// },
// }
// binLogs, _, err := insertCodec.Serialize(999, 888, &insertData)
// assert.Nil(t, err)
// kvs := make(map[string][]byte, len(binLogs))
// paths := make([]string, 0, len(binLogs))
// for i, blob := range binLogs {
// key := path.Join(floatVectorBinlogPath, strconv.Itoa(i))
// paths = append(paths, key)
// kvs[key] = blob.Value[:]
// }
// err = in.chunkManager.MultiWrite(kvs)
// assert.Nil(t, err)
//
// indexMeta := &indexpb.IndexMeta{
// IndexBuildID: indexBuildID1,
// State: commonpb.IndexState_InProgress,
// IndexVersion: 1,
// }
//
// value, err := proto.Marshal(indexMeta)
// assert.Nil(t, err)
// err = in.etcdKV.Save(metaPath1, string(value))
// assert.Nil(t, err)
// req := &indexpb.CreateIndexRequest{
// IndexBuildID: indexBuildID1,
// IndexName: "FloatVector",
// IndexID: indexID,
// Version: 1,
// MetaPath: metaPath1,
// DataPaths: paths,
// TypeParams: []*commonpb.KeyValuePair{
// {
// Key: "dim",
// Value: "8",
// },
// },
// IndexParams: []*commonpb.KeyValuePair{
// {
// Key: "index_type",
// Value: "IVF_SQ8",
// },
// {
// Key: "params",
// Value: "{\"nlist\": 128}",
// },
// {
// Key: "metric_type",
// Value: "L2",
// },
// },
// }
//
// status, err2 := in.CreateIndex(ctx, req)
// assert.Nil(t, err2)
// assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode)
//
// strValue, err3 := in.etcdKV.Load(metaPath1)
// assert.Nil(t, err3)
// indexMetaTmp := indexpb.IndexMeta{}
// err = proto.Unmarshal([]byte(strValue), &indexMetaTmp)
// assert.Nil(t, err)
// for indexMetaTmp.State != commonpb.IndexState_Finished {
// time.Sleep(100 * time.Millisecond)
// strValue, err := in.etcdKV.Load(metaPath1)
// assert.Nil(t, err)
// err = proto.Unmarshal([]byte(strValue), &indexMetaTmp)
// assert.Nil(t, err)
// }
// defer in.chunkManager.MultiRemove(indexMetaTmp.IndexFilePaths)
// defer func() {
// for k := range kvs {
// err = in.chunkManager.Remove(k)
// assert.Nil(t, err)
// }
// }()
//
// defer in.etcdKV.RemoveWithPrefix(metaPath1)
// })
// t.Run("CreateIndex BinaryVector", func(t *testing.T) {
// var insertCodec storage.InsertCodec
//
// insertCodec.Schema = &etcdpb.CollectionMeta{
// ID: collectionID,
// Schema: &schemapb.CollectionSchema{
// Fields: []*schemapb.FieldSchema{
// {
// FieldID: binaryVectorFieldID,
// Name: binaryVectorFieldName,
// IsPrimaryKey: false,
// DataType: schemapb.DataType_BinaryVector,
// },
// },
// },
// }
// data := make(map[UniqueID]storage.FieldData)
// tsData := make([]int64, nb)
// for i := 0; i < nb; i++ {
// tsData[i] = int64(i + 100)
// }
// data[tsFieldID] = &storage.Int64FieldData{
// NumRows: []int64{nb},
// Data: tsData,
// }
// data[binaryVectorFieldID] = &storage.BinaryVectorFieldData{
// NumRows: []int64{nb},
// Data: generateBinaryVectors(),
// Dim: dim,
// }
// insertData := storage.InsertData{
// Data: data,
// Infos: []storage.BlobInfo{
// {
// Length: 10,
// },
// },
// }
// binLogs, _, err := insertCodec.Serialize(999, 888, &insertData)
// assert.Nil(t, err)
// kvs := make(map[string][]byte, len(binLogs))
// paths := make([]string, 0, len(binLogs))
// for i, blob := range binLogs {
// key := path.Join(binaryVectorBinlogPath, strconv.Itoa(i))
// paths = append(paths, key)
// kvs[key] = blob.Value[:]
// }
// err = in.chunkManager.MultiWrite(kvs)
// assert.Nil(t, err)
//
// indexMeta := &indexpb.IndexMeta{
// IndexBuildID: indexBuildID2,
// State: commonpb.IndexState_InProgress,
// IndexVersion: 1,
// }
//
// value, err := proto.Marshal(indexMeta)
// assert.Nil(t, err)
// err = in.etcdKV.Save(metaPath2, string(value))
// assert.Nil(t, err)
// req := &indexpb.CreateIndexRequest{
// IndexBuildID: indexBuildID2,
// IndexName: "BinaryVector",
// IndexID: indexID,
// Version: 1,
// MetaPath: metaPath2,
// DataPaths: paths,
// TypeParams: []*commonpb.KeyValuePair{
// {
// Key: "dim",
// Value: "8",
// },
// },
// IndexParams: []*commonpb.KeyValuePair{
// {
// Key: "index_type",
// Value: "BIN_FLAT",
// },
// {
// Key: "metric_type",
// Value: "JACCARD",
// },
// },
// }
//
// status, err2 := in.CreateIndex(ctx, req)
// assert.Nil(t, err2)
// assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode)
//
// strValue, err3 := in.etcdKV.Load(metaPath2)
// assert.Nil(t, err3)
// indexMetaTmp := indexpb.IndexMeta{}
// err = proto.Unmarshal([]byte(strValue), &indexMetaTmp)
// assert.Nil(t, err)
// for indexMetaTmp.State != commonpb.IndexState_Finished {
// time.Sleep(100 * time.Millisecond)
// strValue, err = in.etcdKV.Load(metaPath2)
// assert.Nil(t, err)
// err = proto.Unmarshal([]byte(strValue), &indexMetaTmp)
// assert.Nil(t, err)
// }
// defer in.chunkManager.MultiRemove(indexMetaTmp.IndexFilePaths)
// defer func() {
// for k := range kvs {
// err = in.chunkManager.Remove(k)
// assert.Nil(t, err)
// }
// }()
//
// defer in.etcdKV.RemoveWithPrefix(metaPath2)
// })
//
// t.Run("Create DeletedIndex", func(t *testing.T) {
// var insertCodec storage.InsertCodec
//
// insertCodec.Schema = &etcdpb.CollectionMeta{
// ID: collectionID,
// Schema: &schemapb.CollectionSchema{
// Fields: []*schemapb.FieldSchema{
// {
// FieldID: floatVectorFieldID,
// Name: floatVectorFieldName,
// IsPrimaryKey: false,
// DataType: schemapb.DataType_FloatVector,
// },
// },
// },
// }
// data := make(map[UniqueID]storage.FieldData)
// tsData := make([]int64, nb)
// for i := 0; i < nb; i++ {
// tsData[i] = int64(i + 100)
// }
// data[tsFieldID] = &storage.Int64FieldData{
// NumRows: []int64{nb},
// Data: tsData,
// }
// data[floatVectorFieldID] = &storage.FloatVectorFieldData{
// NumRows: []int64{nb},
// Data: generateFloatVectors(),
// Dim: dim,
// }
// insertData := storage.InsertData{
// Data: data,
// Infos: []storage.BlobInfo{
// {
// Length: 10,
// },
// },
// }
// binLogs, _, err := insertCodec.Serialize(999, 888, &insertData)
// assert.Nil(t, err)
// kvs := make(map[string][]byte, len(binLogs))
// paths := make([]string, 0, len(binLogs))
// for i, blob := range binLogs {
// key := path.Join(floatVectorBinlogPath, strconv.Itoa(i))
// paths = append(paths, key)
// kvs[key] = blob.Value[:]
// }
// err = in.chunkManager.MultiWrite(kvs)
// assert.Nil(t, err)
//
// indexMeta := &indexpb.IndexMeta{
// IndexBuildID: indexBuildID1,
// State: commonpb.IndexState_InProgress,
// IndexVersion: 1,
// MarkDeleted: true,
// }
//
// value, err := proto.Marshal(indexMeta)
// assert.Nil(t, err)
// err = in.etcdKV.Save(metaPath3, string(value))
// assert.Nil(t, err)
// req := &indexpb.CreateIndexRequest{
// IndexBuildID: indexBuildID1,
// IndexName: "FloatVector",
// IndexID: indexID,
// Version: 1,
// MetaPath: metaPath3,
// DataPaths: paths,
// TypeParams: []*commonpb.KeyValuePair{
// {
// Key: "dim",
// Value: "8",
// },
// },
// IndexParams: []*commonpb.KeyValuePair{
// {
// Key: "index_type",
// Value: "IVF_SQ8",
// },
// {
// Key: "params",
// Value: "{\"nlist\": 128}",
// },
// {
// Key: "metric_type",
// Value: "L2",
// },
// },
// }
//
// status, err2 := in.CreateIndex(ctx, req)
// assert.Nil(t, err2)
// assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode)
// time.Sleep(100 * time.Millisecond)
// strValue, err3 := in.etcdKV.Load(metaPath3)
// assert.Nil(t, err3)
// indexMetaTmp := indexpb.IndexMeta{}
// err = proto.Unmarshal([]byte(strValue), &indexMetaTmp)
// assert.Nil(t, err)
// assert.Equal(t, true, indexMetaTmp.MarkDeleted)
// assert.Equal(t, int64(1), indexMetaTmp.IndexVersion)
// //for indexMetaTmp.State != commonpb.IndexState_Finished {
// // time.Sleep(100 * time.Millisecond)
// // strValue, err := in.etcdKV.Load(metaPath3)
// // assert.Nil(t, err)
// // err = proto.Unmarshal([]byte(strValue), &indexMetaTmp)
// // assert.Nil(t, err)
// //}
// defer in.chunkManager.MultiRemove(indexMetaTmp.IndexFilePaths)
// defer func() {
// for k := range kvs {
// err = in.chunkManager.Remove(k)
// assert.Nil(t, err)
// }
// }()
//
// defer in.etcdKV.RemoveWithPrefix(metaPath3)
// })
//
// t.Run("GetComponentStates", func(t *testing.T) {
// resp, err := in.GetComponentStates(ctx)
// assert.Nil(t, err)
// assert.Equal(t, commonpb.ErrorCode_Success, resp.Status.ErrorCode)
// assert.Equal(t, commonpb.StateCode_Healthy, resp.State.StateCode)
// })
//
// t.Run("GetTimeTickChannel", func(t *testing.T) {
// resp, err := in.GetTimeTickChannel(ctx)
// assert.Nil(t, err)
// assert.Equal(t, commonpb.ErrorCode_Success, resp.Status.ErrorCode)
// })
//
// t.Run("GetStatisticsChannel", func(t *testing.T) {
// resp, err := in.GetStatisticsChannel(ctx)
// assert.Nil(t, err)
// assert.Equal(t, commonpb.ErrorCode_Success, resp.Status.ErrorCode)
// })
//
// t.Run("ShowConfigurations", func(t *testing.T) {
// pattern := "Port"
// req := &internalpb.ShowConfigurationsRequest{
// Base: &commonpb.MsgBase{
// MsgType: commonpb.MsgType_WatchQueryChannels,
// MsgID: rand.Int63(),
// },
// Pattern: pattern,
// }
//
// resp, err := in.ShowConfigurations(ctx, req)
// assert.NoError(t, err)
// assert.Equal(t, commonpb.ErrorCode_Success, resp.Status.ErrorCode)
// assert.Equal(t, 1, len(resp.Configuations))
// assert.Equal(t, "indexnode.port", resp.Configuations[0].Key)
// })
//
// t.Run("GetMetrics_system_info", func(t *testing.T) {
// req, err := metricsinfo.ConstructRequestByMetricType(metricsinfo.SystemInfoMetrics)
// assert.Nil(t, err)
// resp, err := in.GetMetrics(ctx, req)
// assert.Nil(t, err)
// log.Info("GetMetrics_system_info",
// zap.String("resp", resp.Response),
// zap.String("name", resp.ComponentName))
// })
// err = in.etcdKV.RemoveWithPrefix("session/IndexNode")
// assert.Nil(t, err)
//
// resp, err = getEtcdClient().Get(ctx, path.Join(Params.EtcdCfg.MetaRootPath, sessionutil.DefaultServiceRoot, in.session.ServerName))
// assert.Nil(t, err)
// assert.Equal(t, resp.Count, int64(0))
//}
func TestComponentState(t *testing.T) {
var (
factory = &mockFactory{
chunkMgr: &mockChunkmgr{},
}
ctx = context.TODO()
)
Params.Init()
in, err := NewIndexNode(ctx, factory)
assert.Nil(t, err)
in.SetEtcdClient(getEtcdClient())
state, err := in.GetComponentStates(ctx)
assert.Nil(t, err)
assert.Equal(t, state.Status.ErrorCode, commonpb.ErrorCode_Success)
assert.Equal(t, state.State.StateCode, commonpb.StateCode_Abnormal)
assert.Nil(t, in.Init())
state, err = in.GetComponentStates(ctx)
assert.Nil(t, err)
assert.Equal(t, state.Status.ErrorCode, commonpb.ErrorCode_Success)
assert.Equal(t, state.State.StateCode, commonpb.StateCode_Initializing)
assert.Nil(t, in.Start())
state, err = in.GetComponentStates(ctx)
assert.Nil(t, err)
assert.Equal(t, state.Status.ErrorCode, commonpb.ErrorCode_Success)
assert.Equal(t, state.State.StateCode, commonpb.StateCode_Healthy)
assert.Nil(t, in.Stop())
state, err = in.GetComponentStates(ctx)
assert.Nil(t, err)
assert.Equal(t, state.Status.ErrorCode, commonpb.ErrorCode_Success)
assert.Equal(t, state.State.StateCode, commonpb.StateCode_Abnormal)
}
func TestGetTimeTickChannel(t *testing.T) {
var (
factory = &mockFactory{
chunkMgr: &mockChunkmgr{},
}
ctx = context.TODO()
)
Params.Init()
in, err := NewIndexNode(ctx, factory)
assert.Nil(t, err)
ret, err := in.GetTimeTickChannel(ctx)
assert.Nil(t, err)
assert.Equal(t, ret.Status.ErrorCode, commonpb.ErrorCode_Success)
}
func TestGetStatisticChannel(t *testing.T) {
var (
factory = &mockFactory{
chunkMgr: &mockChunkmgr{},
}
ctx = context.TODO()
)
Params.Init()
in, err := NewIndexNode(ctx, factory)
assert.Nil(t, err)
ret, err := in.GetStatisticsChannel(ctx)
assert.Nil(t, err)
assert.Equal(t, ret.Status.ErrorCode, commonpb.ErrorCode_Success)
}
func TestInitErr(t *testing.T) {
// var (
// factory = &mockFactory{}
// ctx = context.TODO()
// )
// in, err := NewIndexNode(ctx, factory)
// assert.Nil(t, err)
// in.SetEtcdClient(getEtcdClient())
// assert.Error(t, in.Init())
}
func setup() {
startEmbedEtcd()
}
func teardown() {
stopEmbedEtcd()
}
func TestMain(m *testing.M) {
setup()
code := m.Run()
teardown()
os.Exit(code)
}