Add unittest for flowgraph (#7512)

internal/util/flowgraph/flow_graph_test.go

@@ -203,7 +203,83 @@ func createExampleFlowGraph() (*TimeTickedFlowGraph, chan float64, chan float64,
 	return fg, inputChan, outputChan, cancel
 }
 
-func TestTimeTickedFlowGraph_Example(t *testing.T) {
+func TestTimeTickedFlowGraphAddNode(t *testing.T) {
+	const MaxQueueLength = 1024
+	inputChan := make(chan float64, MaxQueueLength)
+
+	fg := NewTimeTickedFlowGraph(context.TODO())
+
+	var a Node = &nodeA{
+		BaseNode: BaseNode{
+			maxQueueLength: MaxQueueLength,
+		},
+		inputChan: inputChan,
+	}
+	var b Node = &nodeB{
+		BaseNode: BaseNode{
+			maxQueueLength: MaxQueueLength,
+		},
+	}
+
+	fg.AddNode(a)
+	assert.Equal(t, len(fg.nodeCtx), 1)
+	fg.AddNode(b)
+	assert.Equal(t, len(fg.nodeCtx), 2)
+}
+
+func TestTimeTickedFlowGraphSetEdges(t *testing.T) {
+	const MaxQueueLength = 1024
+	inputChan := make(chan float64, MaxQueueLength)
+
+	fg := NewTimeTickedFlowGraph(context.TODO())
+
+	var a Node = &nodeA{
+		BaseNode: BaseNode{
+			maxQueueLength: MaxQueueLength,
+		},
+		inputChan: inputChan,
+	}
+	var b Node = &nodeB{
+		BaseNode: BaseNode{
+			maxQueueLength: MaxQueueLength,
+		},
+	}
+	var c Node = &nodeC{
+		BaseNode: BaseNode{
+			maxQueueLength: MaxQueueLength,
+		},
+	}
+
+	fg.AddNode(a)
+	fg.AddNode(b)
+	fg.AddNode(c)
+
+	var err = fg.SetEdges(a.Name(),
+		[]string{b.Name()},
+		[]string{c.Name()},
+	)
+	assert.Nil(t, err)
+
+	err = fg.SetEdges("Invalid",
+		[]string{b.Name()},
+		[]string{c.Name()},
+	)
+	assert.Error(t, err)
+
+	err = fg.SetEdges(a.Name(),
+		[]string{"Invalid"},
+		[]string{c.Name()},
+	)
+	assert.Error(t, err)
+
+	err = fg.SetEdges(a.Name(),
+		[]string{b.Name()},
+		[]string{"Invalid"},
+	)
+	assert.Error(t, err)
+}
+
+func TestTimeTickedFlowGraphStart(t *testing.T) {
 	fg, inputChan, outputChan, cancel := createExampleFlowGraph()
 	defer cancel()
 	go fg.Start()
@@ -223,3 +299,9 @@ func TestTimeTickedFlowGraph_Example(t *testing.T) {
 	}()
 	time.Sleep(50 * time.Millisecond)
 }
+
+func TestTimeTickedFlowGraphClose(t *testing.T) {
+	fg, _, _, cancel := createExampleFlowGraph()
+	defer cancel()
+	fg.Close()
+}

internal/util/flowgraph/input_node_test.go

@@ -0,0 +1,80 @@
+// Copyright (C) 2019-2020 Zilliz. All rights reserved.
+//
+// Licensed 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 flowgraph
+
+import (
+	"context"
+	"os"
+	"sync"
+	"testing"
+
+	"github.com/milvus-io/milvus/internal/msgstream"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestInputNode(t *testing.T) {
+	os.Setenv("ROCKSMQ_PATH", "/tmp/MilvusTest/FlowGraph/TestInputNode")
+	msFactory := msgstream.NewRmsFactory()
+	m := map[string]interface{}{}
+	err := msFactory.SetParams(m)
+	assert.Nil(t, err)
+
+	msgStream, _ := msFactory.NewMsgStream(context.TODO())
+	channels := []string{"cc"}
+	msgStream.AsConsumer(channels, "sub")
+	msgStream.Start()
+
+	msgPack := generateMsgPack()
+	produceStream, _ := msFactory.NewMsgStream(context.TODO())
+	produceStream.AsProducer(channels)
+	produceStream.Produce(&msgPack)
+
+	nodeName := "input_node"
+	inputNode := &InputNode{
+		inStream: &msgStream,
+		name:     nodeName,
+	}
+	inputNode.Close()
+
+	isInputNode := inputNode.IsInputNode()
+	assert.True(t, isInputNode)
+
+	name := inputNode.Name()
+	assert.Equal(t, name, nodeName)
+
+	stream := inputNode.InStream()
+	assert.NotNil(t, stream)
+
+	var waitGroup sync.WaitGroup
+	OperateFunc := func() {
+		msgs := make([]Msg, 0)
+		output := inputNode.Operate(msgs)
+		assert.Greater(t, len(output), 0)
+		msgStream.Close()
+		waitGroup.Done()
+	}
+
+	waitGroup.Add(1)
+	go OperateFunc()
+	waitGroup.Wait()
+}
+
+func TestNewInputNode(t *testing.T) {
+	nodeName := "input_node"
+	var maxQueueLength int32 = 0
+	var maxParallelism int32 = 100
+	node := NewInputNode(nil, nodeName, maxQueueLength, maxParallelism)
+	assert.NotNil(t, node)
+	assert.Equal(t, node.name, nodeName)
+	assert.Equal(t, node.maxQueueLength, maxQueueLength)
+	assert.Equal(t, node.maxParallelism, maxParallelism)
+}

internal/util/flowgraph/message_test.go

@@ -0,0 +1,127 @@
+// Copyright (C) 2019-2020 Zilliz. All rights reserved.
+//
+// Licensed 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 flowgraph
+
+import (
+	"context"
+	"testing"
+
+	"github.com/milvus-io/milvus/internal/msgstream"
+	"github.com/stretchr/testify/assert"
+)
+
+type MockMsg struct {
+	Ctx context.Context
+}
+
+func (bm *MockMsg) TraceCtx() context.Context {
+	return bm.Ctx
+}
+
+func (bm *MockMsg) SetTraceCtx(ctx context.Context) {
+	bm.Ctx = ctx
+}
+
+func (bm *MockMsg) ID() msgstream.UniqueID {
+	return 0
+}
+
+func (bm *MockMsg) BeginTs() Timestamp {
+	return 0
+}
+
+func (bm *MockMsg) EndTs() Timestamp {
+	return 0
+}
+
+func (bm *MockMsg) Type() msgstream.MsgType {
+	return 0
+}
+
+func (bm *MockMsg) SourceID() int64 {
+	return 0
+}
+
+func (bm *MockMsg) HashKeys() []uint32 {
+	return []uint32{0}
+}
+
+func (bm *MockMsg) Marshal(msgstream.TsMsg) (msgstream.MarshalType, error) {
+	return nil, nil
+}
+
+func (bm *MockMsg) Unmarshal(msgstream.MarshalType) (msgstream.TsMsg, error) {
+	return nil, nil
+}
+
+func (bm *MockMsg) Position() *MsgPosition {
+	return nil
+}
+
+func (bm *MockMsg) SetPosition(position *MsgPosition) {
+
+}
+
+func TestGenerateMsgStreamMsg(t *testing.T) {
+	messages := make([]msgstream.TsMsg, 1)
+	messages[0] = &MockMsg{
+		Ctx: context.TODO(),
+	}
+
+	timestampMin := uint64(0)
+	timestampMax := uint64(1000)
+	streamMsg := GenerateMsgStreamMsg(messages, timestampMin, timestampMax, nil, nil)
+	assert.NotNil(t, streamMsg)
+	assert.Equal(t, len(streamMsg.tsMessages), len(messages))
+	assert.Equal(t, streamMsg.timestampMin, timestampMin)
+	assert.Equal(t, streamMsg.timestampMax, timestampMax)
+	assert.Nil(t, streamMsg.startPositions)
+	assert.Nil(t, streamMsg.endPositions)
+}
+
+func TestMsgStreamMsg(t *testing.T) {
+	messages := make([]msgstream.TsMsg, 1)
+	messages[0] = &MockMsg{
+		Ctx: context.TODO(),
+	}
+
+	var timestampMin Timestamp = 0
+	var timestampMax Timestamp = 100
+	streamMsg := &MsgStreamMsg{
+		tsMessages:     messages,
+		timestampMin:   timestampMin,
+		timestampMax:   timestampMax,
+		startPositions: nil,
+		endPositions:   nil,
+	}
+
+	tt := streamMsg.TimeTick()
+	assert.Equal(t, tt, timestampMax)
+
+	nodeID := streamMsg.DownStreamNodeIdx()
+	assert.Equal(t, nodeID, 0)
+
+	msgs := streamMsg.TsMessages()
+	assert.Equal(t, len(msgs), len(messages))
+
+	min := streamMsg.TimestampMin()
+	assert.Equal(t, min, timestampMin)
+
+	max := streamMsg.TimestampMax()
+	assert.Equal(t, max, timestampMax)
+
+	startPos := streamMsg.StartPositions()
+	assert.Nil(t, startPos)
+
+	endPos := streamMsg.EndPositions()
+	assert.Nil(t, endPos)
+}

internal/util/flowgraph/node_test.go

@@ -0,0 +1,120 @@
+// Copyright (C) 2019-2020 Zilliz. All rights reserved.
+//
+// Licensed 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 flowgraph
+
+import (
+	"context"
+	"math"
+	"os"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/milvus-io/milvus/internal/msgstream"
+	"github.com/milvus-io/milvus/internal/proto/commonpb"
+	"github.com/milvus-io/milvus/internal/proto/internalpb"
+	"github.com/stretchr/testify/assert"
+)
+
+func generateMsgPack() msgstream.MsgPack {
+	msgPack := msgstream.MsgPack{}
+	baseMsg := msgstream.BaseMsg{
+		BeginTimestamp: uint64(time.Now().Unix()),
+		EndTimestamp:   uint64(time.Now().Unix() + 1),
+		HashValues:     []uint32{0},
+	}
+	timeTickResult := internalpb.TimeTickMsg{
+		Base: &commonpb.MsgBase{
+			MsgType:   commonpb.MsgType_TimeTick,
+			MsgID:     0,
+			Timestamp: math.MaxUint64,
+			SourceID:  0,
+		},
+	}
+	timeTickMsg := &msgstream.TimeTickMsg{
+		BaseMsg:     baseMsg,
+		TimeTickMsg: timeTickResult,
+	}
+	msgPack.Msgs = append(msgPack.Msgs, timeTickMsg)
+
+	return msgPack
+}
+
+func TestNodeStart(t *testing.T) {
+	os.Setenv("ROCKSMQ_PATH", "/tmp/MilvusTest/FlowGraph/TestNodeStart")
+	msFactory := msgstream.NewRmsFactory()
+	m := map[string]interface{}{}
+	err := msFactory.SetParams(m)
+	assert.Nil(t, err)
+
+	msgStream, _ := msFactory.NewMsgStream(context.TODO())
+	channels := []string{"cc"}
+	msgStream.AsConsumer(channels, "sub")
+
+	produceStream, _ := msFactory.NewMsgStream(context.TODO())
+	produceStream.AsProducer(channels)
+
+	msgPack := generateMsgPack()
+	produceStream.Produce(&msgPack)
+	time.Sleep(time.Millisecond * 2)
+	msgPack = generateMsgPack()
+	produceStream.Produce(&msgPack)
+
+	nodeName := "input_node"
+	inputNode := &InputNode{
+		inStream: &msgStream,
+		name:     nodeName,
+	}
+
+	node := &nodeCtx{
+		node:                   inputNode,
+		inputChannels:          make([]chan Msg, 2),
+		downstreamInputChanIdx: make(map[string]int),
+	}
+
+	for i := 0; i < len(node.inputChannels); i++ {
+		node.inputChannels[i] = make(chan Msg)
+	}
+
+	ctx, cancel := context.WithCancel(context.TODO())
+	defer cancel()
+
+	var waitGroup sync.WaitGroup
+	waitGroup.Add(1)
+	go node.Start(ctx, &waitGroup)
+
+	node.Close()
+}
+
+func TestBaseNode(t *testing.T) {
+	node := &BaseNode{
+		maxQueueLength: 10,
+		maxParallelism: 50,
+	}
+
+	x := node.MaxQueueLength()
+	assert.Equal(t, x, node.maxQueueLength)
+
+	x = node.MaxParallelism()
+	assert.Equal(t, x, node.maxParallelism)
+
+	var val int32 = 3
+	node.SetMaxQueueLength(val)
+	assert.Equal(t, val, node.maxQueueLength)
+
+	node.SetMaxParallelism(val)
+	assert.Equal(t, val, node.maxParallelism)
+
+	assert.Equal(t, false, node.IsInputNode())
+
+	node.Close()
+}