mirror of https://github.com/milvus-io/milvus.git
298 lines
9.8 KiB
Go
298 lines
9.8 KiB
Go
package master
|
|
|
|
import (
|
|
"context"
|
|
"log"
|
|
"math/rand"
|
|
"os"
|
|
"strconv"
|
|
"testing"
|
|
|
|
"github.com/zilliztech/milvus-distributed/internal/util/tsoutil"
|
|
|
|
"github.com/golang/protobuf/proto"
|
|
"github.com/stretchr/testify/assert"
|
|
ms "github.com/zilliztech/milvus-distributed/internal/msgstream"
|
|
"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
|
|
"github.com/zilliztech/milvus-distributed/internal/proto/internalpb"
|
|
internalPb "github.com/zilliztech/milvus-distributed/internal/proto/internalpb"
|
|
"github.com/zilliztech/milvus-distributed/internal/proto/masterpb"
|
|
"github.com/zilliztech/milvus-distributed/internal/proto/schemapb"
|
|
"github.com/zilliztech/milvus-distributed/internal/proto/servicepb"
|
|
"go.uber.org/zap"
|
|
"google.golang.org/grpc"
|
|
)
|
|
|
|
var testPORT = 53200
|
|
|
|
func genMasterTestPort() int64 {
|
|
testPORT++
|
|
return int64(testPORT)
|
|
}
|
|
|
|
func refreshMasterAddress() {
|
|
masterPort := genMasterTestPort()
|
|
Params.Port = int(masterPort)
|
|
masterAddr := makeMasterAddress(masterPort)
|
|
Params.Address = masterAddr
|
|
}
|
|
|
|
func makeMasterAddress(port int64) string {
|
|
masterAddr := "127.0.0.1:" + strconv.FormatInt(port, 10)
|
|
return masterAddr
|
|
}
|
|
|
|
func makeNewChannalNames(names []string, suffix string) []string {
|
|
var ret []string
|
|
for _, name := range names {
|
|
ret = append(ret, name+suffix)
|
|
}
|
|
return ret
|
|
}
|
|
|
|
func refreshChannelNames() {
|
|
suffix := "_test" + strconv.FormatInt(rand.Int63n(100), 10)
|
|
Params.DDChannelNames = makeNewChannalNames(Params.DDChannelNames, suffix)
|
|
Params.WriteNodeTimeTickChannelNames = makeNewChannalNames(Params.WriteNodeTimeTickChannelNames, suffix)
|
|
Params.InsertChannelNames = makeNewChannalNames(Params.InsertChannelNames, suffix)
|
|
Params.K2SChannelNames = makeNewChannalNames(Params.K2SChannelNames, suffix)
|
|
Params.ProxyTimeTickChannelNames = makeNewChannalNames(Params.ProxyTimeTickChannelNames, suffix)
|
|
}
|
|
|
|
func receiveTimeTickMsg(stream *ms.MsgStream) bool {
|
|
for {
|
|
result := (*stream).Consume()
|
|
if len(result.Msgs) > 0 {
|
|
return true
|
|
}
|
|
}
|
|
}
|
|
|
|
func getTimeTickMsgPack(ttmsgs [][2]uint64) *ms.MsgPack {
|
|
msgPack := ms.MsgPack{}
|
|
for _, vi := range ttmsgs {
|
|
msgPack.Msgs = append(msgPack.Msgs, getTtMsg(internalPb.MsgType_kTimeTick, UniqueID(vi[0]), Timestamp(vi[1])))
|
|
}
|
|
return &msgPack
|
|
}
|
|
|
|
func TestMain(m *testing.M) {
|
|
Init()
|
|
refreshMasterAddress()
|
|
refreshChannelNames()
|
|
etcdAddr := Params.EtcdAddress
|
|
gTestTsoAllocator = NewGlobalTSOAllocator("timestamp", tsoutil.NewTSOKVBase([]string{etcdAddr}, "/test/root/kv", "tso"))
|
|
gTestIDAllocator = NewGlobalIDAllocator("idTimestamp", tsoutil.NewTSOKVBase([]string{etcdAddr}, "/test/root/kv", "gid"))
|
|
exitCode := m.Run()
|
|
os.Exit(exitCode)
|
|
}
|
|
|
|
func TestMaster(t *testing.T) {
|
|
Init()
|
|
refreshMasterAddress()
|
|
pulsarAddr := Params.PulsarAddress
|
|
Params.ProxyIDList = []UniqueID{0}
|
|
//Param
|
|
// Creates server.
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
svr, err := CreateServer(ctx)
|
|
if err != nil {
|
|
log.Print("create server failed", zap.Error(err))
|
|
}
|
|
|
|
if err := svr.Run(int64(Params.Port)); err != nil {
|
|
log.Fatal("run server failed", zap.Error(err))
|
|
}
|
|
|
|
proxyTimeTickStream := ms.NewPulsarMsgStream(ctx, 1024) //input stream
|
|
proxyTimeTickStream.SetPulsarClient(pulsarAddr)
|
|
proxyTimeTickStream.CreatePulsarProducers(Params.ProxyTimeTickChannelNames)
|
|
proxyTimeTickStream.Start()
|
|
|
|
writeNodeStream := ms.NewPulsarMsgStream(ctx, 1024) //input stream
|
|
writeNodeStream.SetPulsarClient(pulsarAddr)
|
|
writeNodeStream.CreatePulsarProducers(Params.WriteNodeTimeTickChannelNames)
|
|
writeNodeStream.Start()
|
|
|
|
ddMs := ms.NewPulsarMsgStream(ctx, 1024)
|
|
ddMs.SetPulsarClient(pulsarAddr)
|
|
ddMs.CreatePulsarConsumers(Params.DDChannelNames, "DDStream", ms.NewUnmarshalDispatcher(), 1024)
|
|
ddMs.Start()
|
|
|
|
dMMs := ms.NewPulsarMsgStream(ctx, 1024)
|
|
dMMs.SetPulsarClient(pulsarAddr)
|
|
dMMs.CreatePulsarConsumers(Params.InsertChannelNames, "DMStream", ms.NewUnmarshalDispatcher(), 1024)
|
|
dMMs.Start()
|
|
|
|
k2sMs := ms.NewPulsarMsgStream(ctx, 1024)
|
|
k2sMs.SetPulsarClient(pulsarAddr)
|
|
k2sMs.CreatePulsarConsumers(Params.K2SChannelNames, "K2SStream", ms.NewUnmarshalDispatcher(), 1024)
|
|
k2sMs.Start()
|
|
|
|
ttsoftmsgs := [][2]uint64{
|
|
{0, 10},
|
|
}
|
|
msgSoftPackAddr := getTimeTickMsgPack(ttsoftmsgs)
|
|
|
|
proxyTimeTickStream.Produce(msgSoftPackAddr)
|
|
var dMMsgstream ms.MsgStream = dMMs
|
|
assert.True(t, receiveTimeTickMsg(&dMMsgstream))
|
|
var ddMsgstream ms.MsgStream = ddMs
|
|
assert.True(t, receiveTimeTickMsg(&ddMsgstream))
|
|
|
|
tthardmsgs := [][2]int{
|
|
{3, 10},
|
|
}
|
|
|
|
msghardPackAddr := getMsgPack(tthardmsgs)
|
|
writeNodeStream.Produce(msghardPackAddr)
|
|
var k2sMsgstream ms.MsgStream = k2sMs
|
|
assert.True(t, receiveTimeTickMsg(&k2sMsgstream))
|
|
|
|
conn, err := grpc.DialContext(ctx, Params.Address, grpc.WithInsecure(), grpc.WithBlock())
|
|
assert.Nil(t, err)
|
|
defer conn.Close()
|
|
|
|
cli := masterpb.NewMasterClient(conn)
|
|
|
|
sch := schemapb.CollectionSchema{
|
|
Name: "name" + strconv.FormatUint(rand.Uint64(), 10),
|
|
Description: "test collection",
|
|
AutoID: false,
|
|
Fields: []*schemapb.FieldSchema{},
|
|
}
|
|
|
|
schemaBytes, err := proto.Marshal(&sch)
|
|
assert.Nil(t, err)
|
|
|
|
createCollectionReq := internalpb.CreateCollectionRequest{
|
|
MsgType: internalpb.MsgType_kCreateCollection,
|
|
ReqID: 1,
|
|
Timestamp: 11,
|
|
ProxyID: 1,
|
|
Schema: &commonpb.Blob{Value: schemaBytes},
|
|
}
|
|
st, err := cli.CreateCollection(ctx, &createCollectionReq)
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, st.ErrorCode, commonpb.ErrorCode_SUCCESS)
|
|
|
|
var consumeMsg ms.MsgStream = ddMs
|
|
var createCollectionMsg *ms.CreateCollectionMsg
|
|
for {
|
|
result := consumeMsg.Consume()
|
|
if len(result.Msgs) > 0 {
|
|
msgs := result.Msgs
|
|
for _, v := range msgs {
|
|
createCollectionMsg = v.(*ms.CreateCollectionMsg)
|
|
}
|
|
break
|
|
}
|
|
}
|
|
assert.Equal(t, createCollectionReq.MsgType, createCollectionMsg.CreateCollectionRequest.MsgType)
|
|
assert.Equal(t, createCollectionReq.ReqID, createCollectionMsg.CreateCollectionRequest.ReqID)
|
|
assert.Equal(t, createCollectionReq.Timestamp, createCollectionMsg.CreateCollectionRequest.Timestamp)
|
|
assert.Equal(t, createCollectionReq.ProxyID, createCollectionMsg.CreateCollectionRequest.ProxyID)
|
|
assert.Equal(t, createCollectionReq.Schema.Value, createCollectionMsg.CreateCollectionRequest.Schema.Value)
|
|
|
|
////////////////////////////CreatePartition////////////////////////
|
|
partitionName := "partitionName" + strconv.FormatUint(rand.Uint64(), 10)
|
|
createPartitionReq := internalpb.CreatePartitionRequest{
|
|
MsgType: internalpb.MsgType_kCreatePartition,
|
|
ReqID: 1,
|
|
Timestamp: 11,
|
|
ProxyID: 1,
|
|
PartitionName: &servicepb.PartitionName{
|
|
CollectionName: sch.Name,
|
|
Tag: partitionName,
|
|
},
|
|
}
|
|
|
|
st, err = cli.CreatePartition(ctx, &createPartitionReq)
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, st.ErrorCode, commonpb.ErrorCode_SUCCESS)
|
|
|
|
var createPartitionMsg *ms.CreatePartitionMsg
|
|
for {
|
|
result := consumeMsg.Consume()
|
|
if len(result.Msgs) > 0 {
|
|
msgs := result.Msgs
|
|
for _, v := range msgs {
|
|
createPartitionMsg = v.(*ms.CreatePartitionMsg)
|
|
}
|
|
break
|
|
}
|
|
}
|
|
assert.Equal(t, createPartitionReq.MsgType, createPartitionMsg.CreatePartitionRequest.MsgType)
|
|
assert.Equal(t, createPartitionReq.ReqID, createPartitionMsg.CreatePartitionRequest.ReqID)
|
|
assert.Equal(t, createPartitionReq.Timestamp, createPartitionMsg.CreatePartitionRequest.Timestamp)
|
|
assert.Equal(t, createPartitionReq.ProxyID, createPartitionMsg.CreatePartitionRequest.ProxyID)
|
|
assert.Equal(t, createPartitionReq.PartitionName.CollectionName, createPartitionMsg.CreatePartitionRequest.PartitionName.CollectionName)
|
|
assert.Equal(t, createPartitionReq.PartitionName.Tag, createPartitionMsg.CreatePartitionRequest.PartitionName.Tag)
|
|
|
|
////////////////////////////DropPartition////////////////////////
|
|
dropPartitionReq := internalpb.DropPartitionRequest{
|
|
MsgType: internalpb.MsgType_kDropPartition,
|
|
ReqID: 1,
|
|
Timestamp: 11,
|
|
ProxyID: 1,
|
|
PartitionName: &servicepb.PartitionName{
|
|
CollectionName: sch.Name,
|
|
Tag: partitionName,
|
|
},
|
|
}
|
|
|
|
st, err = cli.DropPartition(ctx, &dropPartitionReq)
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, st.ErrorCode, commonpb.ErrorCode_SUCCESS)
|
|
|
|
var dropPartitionMsg *ms.DropPartitionMsg
|
|
for {
|
|
result := consumeMsg.Consume()
|
|
if len(result.Msgs) > 0 {
|
|
msgs := result.Msgs
|
|
for _, v := range msgs {
|
|
dropPartitionMsg = v.(*ms.DropPartitionMsg)
|
|
}
|
|
break
|
|
}
|
|
}
|
|
assert.Equal(t, dropPartitionReq.MsgType, dropPartitionMsg.DropPartitionRequest.MsgType)
|
|
assert.Equal(t, dropPartitionReq.ReqID, dropPartitionMsg.DropPartitionRequest.ReqID)
|
|
assert.Equal(t, dropPartitionReq.Timestamp, dropPartitionMsg.DropPartitionRequest.Timestamp)
|
|
assert.Equal(t, dropPartitionReq.ProxyID, dropPartitionMsg.DropPartitionRequest.ProxyID)
|
|
assert.Equal(t, dropPartitionReq.PartitionName.CollectionName, dropPartitionMsg.DropPartitionRequest.PartitionName.CollectionName)
|
|
|
|
////////////////////////////DropCollection////////////////////////
|
|
dropCollectionReq := internalpb.DropCollectionRequest{
|
|
MsgType: internalpb.MsgType_kDropCollection,
|
|
ReqID: 1,
|
|
Timestamp: 11,
|
|
ProxyID: 1,
|
|
CollectionName: &servicepb.CollectionName{CollectionName: sch.Name},
|
|
}
|
|
|
|
st, err = cli.DropCollection(ctx, &dropCollectionReq)
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, st.ErrorCode, commonpb.ErrorCode_SUCCESS)
|
|
|
|
var dropCollectionMsg *ms.DropCollectionMsg
|
|
for {
|
|
result := consumeMsg.Consume()
|
|
if len(result.Msgs) > 0 {
|
|
msgs := result.Msgs
|
|
for _, v := range msgs {
|
|
dropCollectionMsg = v.(*ms.DropCollectionMsg)
|
|
}
|
|
break
|
|
}
|
|
}
|
|
assert.Equal(t, dropCollectionReq.MsgType, dropCollectionMsg.DropCollectionRequest.MsgType)
|
|
assert.Equal(t, dropCollectionReq.ReqID, dropCollectionMsg.DropCollectionRequest.ReqID)
|
|
assert.Equal(t, dropCollectionReq.Timestamp, dropCollectionMsg.DropCollectionRequest.Timestamp)
|
|
assert.Equal(t, dropCollectionReq.ProxyID, dropCollectionMsg.DropCollectionRequest.ProxyID)
|
|
assert.Equal(t, dropCollectionReq.CollectionName.CollectionName, dropCollectionMsg.DropCollectionRequest.CollectionName.CollectionName)
|
|
|
|
cancel()
|
|
svr.Close()
|
|
}
|