milvus/internal/indexnode/indexnode_test.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-proto/go-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.IndexFileKeys)
//		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.IndexFileKeys)
//		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.IndexFileKeys)
//		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)
}