milvus/internal/datanode/data_node_test.go

package datanode

import (
	"context"
	"fmt"
	"log"
	"math/rand"
	"os"
	"strconv"
	"testing"
	"time"

	// "github.com/stretchr/testify/assert"
	// "github.com/stretchr/testify/require"
	"go.etcd.io/etcd/clientv3"
	"go.uber.org/zap"

	"github.com/zilliztech/milvus-distributed/internal/master"
)

func makeNewChannelNames(names []string, suffix string) []string {
	var ret []string
	for _, name := range names {
		ret = append(ret, name+suffix)
	}
	return ret
}

func refreshChannelNames() {
	suffix := "-test-data-node" + strconv.FormatInt(rand.Int63n(100), 10)
	Params.DDChannelNames = makeNewChannelNames(Params.DDChannelNames, suffix)
	Params.InsertChannelNames = makeNewChannelNames(Params.InsertChannelNames, suffix)
}

func startMaster(ctx context.Context) {
	master.Init()
	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)
	}

	masterPort := 53101
	master.Params.Port = masterPort
	svr, err := master.CreateServer(ctx)
	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())
	Params.MasterAddress = master.Params.Address + ":" + strconv.Itoa(masterPort)
}

func TestMain(m *testing.M) {
	Params.Init()
	refreshChannelNames()
	const ctxTimeInMillisecond = 2000
	const closeWithDeadline = true
	var ctx context.Context

	if closeWithDeadline {
		var cancel context.CancelFunc
		d := time.Now().Add(ctxTimeInMillisecond * time.Millisecond)
		ctx, cancel = context.WithDeadline(context.Background(), d)
		defer cancel()
	} else {
		ctx = context.Background()
	}

	startMaster(ctx)
	// p := Params
	// fmt.Println(p)
	exitCode := m.Run()
	os.Exit(exitCode)
}

func newDataNode() *DataNode {

	const ctxTimeInMillisecond = 2000
	const closeWithDeadline = true
	var ctx context.Context

	if closeWithDeadline {
		var cancel context.CancelFunc
		d := time.Now().Add(ctxTimeInMillisecond * time.Millisecond)
		ctx, cancel = context.WithDeadline(context.Background(), d)
		go func() {
			<-ctx.Done()
			cancel()
		}()
	} else {
		ctx = context.Background()
	}

	svr := NewDataNode(ctx, 0)
	return svr

}

// func genTestCollectionMeta(collectionName string, collectionID UniqueID, isBinary bool) *etcdpb.CollectionMeta {
//     var fieldVec schemapb.FieldSchema
//     if isBinary {
//         fieldVec = schemapb.FieldSchema{
//             FieldID:      UniqueID(100),
//             Name:         "vec",
//             IsPrimaryKey: false,
//             DataType:     schemapb.DataType_VECTOR_BINARY,
//             TypeParams: []*commonpb.KeyValuePair{
//                 {
//                     Key:   "dim",
//                     Value: "128",
//                 },
//             },
//             IndexParams: []*commonpb.KeyValuePair{
//                 {
//                     Key:   "metric_type",
//                     Value: "JACCARD",
//                 },
//             },
//         }
//     } else {
//         fieldVec = schemapb.FieldSchema{
//             FieldID:      UniqueID(100),
//             Name:         "vec",
//             IsPrimaryKey: false,
//             DataType:     schemapb.DataType_VECTOR_FLOAT,
//             TypeParams: []*commonpb.KeyValuePair{
//                 {
//                     Key:   "dim",
//                     Value: "16",
//                 },
//             },
//             IndexParams: []*commonpb.KeyValuePair{
//                 {
//                     Key:   "metric_type",
//                     Value: "L2",
//                 },
//             },
//         }
//     }
//
//     fieldInt := schemapb.FieldSchema{
//         FieldID:      UniqueID(101),
//         Name:         "age",
//         IsPrimaryKey: false,
//         DataType:     schemapb.DataType_INT32,
//     }
//
//     schema := schemapb.CollectionSchema{
//         Name:   collectionName,
//         AutoID: true,
//         Fields: []*schemapb.FieldSchema{
//             &fieldVec, &fieldInt,
//         },
//     }
//
//     collectionMeta := etcdpb.CollectionMeta{
//         ID:            collectionID,
//         Schema:        &schema,
//         CreateTime:    Timestamp(0),
//         SegmentIDs:    []UniqueID{0},
//         PartitionTags: []string{"default"},
//     }
//
//     return &collectionMeta
// }

// func initTestMeta(t *testing.T, node *DataNode, collectionName string, collectionID UniqueID, segmentID UniqueID, optional ...bool) {
//     isBinary := false
//     if len(optional) > 0 {
//         isBinary = optional[0]
//     }
//     collectionMeta := genTestCollectionMeta(collectionName, collectionID, isBinary)
//
//     schemaBlob := proto.MarshalTextString(collectionMeta.Schema)
//     require.NotEqual(t, "", schemaBlob)
//
//     var err = node.replica.addCollection(collectionMeta.ID, schemaBlob)
//     require.NoError(t, err)
//
//     collection, err := node.replica.getCollectionByName(collectionName)
//     require.NoError(t, err)
//     require.Equal(t, collection.Name(), collectionName)
//     require.Equal(t, collection.ID(), collectionID)
//     require.Equal(t, node.replica.getCollectionNum(), 1)
//
//     err = node.replica.addPartition(collection.ID(), collectionMeta.PartitionTags[0])
//     require.NoError(t, err)
//
//     err = node.replica.addSegment(segmentID, collectionMeta.PartitionTags[0], collectionID)
//     require.NoError(t, err)
// }