indexnode_test.go

package indexnode

import (
	"context"
	"fmt"
	"log"
	"os"
	"strconv"
	"testing"
	"time"

	"go.etcd.io/etcd/clientv3"

	"go.uber.org/zap"

	"github.com/stretchr/testify/assert"
	indexnodeclient "github.com/zilliztech/milvus-distributed/internal/indexnode/client"
	"github.com/zilliztech/milvus-distributed/internal/master"
	"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
)

var ctx context.Context
var cancel func()

var buildClient *indexnodeclient.Client

var builderServer *Builder

var masterPort = 53101
var masterServer *master.Master

func makeMasterAddress(port int64) string {
	masterAddr := "127.0.0.1:" + strconv.FormatInt(port, 10)
	return masterAddr
}

func refreshMasterAddress() {
	masterAddr := makeMasterAddress(int64(masterPort))
	Params.MasterAddress = masterAddr
	master.Params.Port = masterPort
}

func startMaster(ctx context.Context) {
	master.Init()
	refreshMasterAddress()
	etcdAddr := master.Params.EtcdAddress
	metaRootPath := master.Params.MetaRootPath

	etcdCli, err := clientv3.New(clientv3.Config{Endpoints: []string{etcdAddr}})
	if err != nil {
		panic(err)
	}
	_, err = etcdCli.Delete(context.TODO(), metaRootPath, clientv3.WithPrefix())
	if err != nil {
		panic(err)
	}

	svr, err := master.CreateServer(ctx)
	masterServer = svr
	if err != nil {
		log.Print("create server failed", zap.Error(err))
	}
	if err := svr.Run(int64(master.Params.Port)); err != nil {
		log.Fatal("run server failed", zap.Error(err))
	}

	fmt.Println("Waiting for server!", svr.IsServing())

}

func startBuilder(ctx context.Context) {
	var err error
	builderServer, err = CreateBuilder(ctx)
	if err != nil {
		log.Print("create builder failed", zap.Error(err))
	}

	// TODO: change to wait until master is ready
	if err := builderServer.Start(); err != nil {
		log.Fatal("run builder failed", zap.Error(err))
	}
}

func setup() {
	Params.Init()
	ctx, cancel = context.WithCancel(context.Background())
	startMaster(ctx)
	startBuilder(ctx)
	addr := Params.Address
	var err error
	buildClient, err = indexnodeclient.NewBuildIndexClient(ctx, addr)
	if err != nil {
		panic("Create buildClient Failed!")
	}

}

func shutdown() {
	cancel()
	builderServer.Close()
	masterServer.Close()
}

func TestMain(m *testing.M) {
	setup()
	code := m.Run()
	shutdown()
	os.Exit(code)
}

func TestBuilder_GRPC(t *testing.T) {
	typeParams := make(map[string]string)
	typeParams["a"] = "1"
	indexParams := make(map[string]string)
	indexParams["b"] = "2"
	columnDataPaths := []string{"dataA", "dataB"}
	indexID, err := buildClient.BuildIndex(columnDataPaths, typeParams, indexParams)
	assert.Nil(t, err)

	time.Sleep(time.Second * 3)

	description, err := buildClient.GetIndexStates([]UniqueID{indexID})
	assert.Nil(t, err)
	assert.Equal(t, commonpb.IndexState_INPROGRESS, description.States[0].State)
	assert.Equal(t, indexID, description.States[0].IndexID)

	indexDataPaths, err := buildClient.GetIndexFilePaths([]UniqueID{indexID})
	assert.Nil(t, err)
	assert.Nil(t, indexDataPaths[0])
}