Add datacoord unit tests (#7197)

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>

internal/datacoord/cluster.go

@@ -90,7 +90,7 @@ func withAssignPolicy(p channelAssignPolicy) ClusterOption {
 }
 
 func defaultRegisterPolicy() dataNodeRegisterPolicy {
-	return newAssiggBufferRegisterPolicy()
+	return newAssignBufferRegisterPolicy()
 }
 
 func defaultUnregisterPolicy() dataNodeUnregisterPolicy {
@@ -257,7 +257,7 @@ func (c *Cluster) handleEvent(node *NodeInfo) {
 				resp, err := cli.FlushSegments(tCtx, req)
 				cancel()
 				if err = VerifyResponse(resp, err); err != nil {
-					log.Warn("failed to flush segments", zap.String("addr", node.Info.GetAddress()))
+					log.Warn("failed to flush segments", zap.String("addr", node.Info.GetAddress()), zap.Error(err))
 				}
 			default:
 				log.Warn("unknown event type", zap.Any("type", event.Type))

internal/datacoord/policy.go

@@ -39,7 +39,7 @@ func newEmptyRegisterPolicy() dataNodeRegisterPolicy {
 	return emptyRegister
 }
 
-func newAssiggBufferRegisterPolicy() dataNodeRegisterPolicy {
+func newAssignBufferRegisterPolicy() dataNodeRegisterPolicy {
 	return registerAssignWithBuffer
 }
 
@@ -102,33 +102,6 @@ func newEmptyUnregisterPolicy() dataNodeUnregisterPolicy {
 // channelAssignFunc, function shortcut for policy
 type channelAssignPolicy func(cluster []*NodeInfo, channel string, collectionID UniqueID) []*NodeInfo
 
-// deprecated
-// test logic, assign channel to all existing data node, works fine only when there is only one data node!
-var assignAllFunc channelAssignPolicy = func(cluster []*NodeInfo, channel string, collectionID UniqueID) []*NodeInfo {
-	ret := make([]*NodeInfo, 0)
-	for _, node := range cluster {
-		has := false
-		for _, ch := range node.Info.GetChannels() {
-			if ch.Name == channel {
-				has = true
-				break
-			}
-		}
-		if has {
-			continue
-		}
-		c := &datapb.ChannelStatus{
-			Name:         channel,
-			State:        datapb.ChannelWatchState_Uncomplete,
-			CollectionID: collectionID,
-		}
-		n := node.Clone(AddChannels([]*datapb.ChannelStatus{c}))
-		ret = append(ret, n)
-	}
-
-	return ret
-}
-
 // balanced assign channel, select the datanode with least amount of channels to assign
 var balancedAssignFunc channelAssignPolicy = func(cluster []*NodeInfo, channel string, collectionID UniqueID) []*NodeInfo {
 	if len(cluster) == 0 {
@@ -161,10 +134,6 @@ var balancedAssignFunc channelAssignPolicy = func(cluster []*NodeInfo, channel s
 	return ret
 }
 
-func newAssignAllPolicy() channelAssignPolicy {
-	return assignAllFunc
-}
-
 func newBalancedAssignPolicy() channelAssignPolicy {
 	return balancedAssignFunc
 }

internal/datacoord/policy_test.go

@@ -78,3 +78,94 @@ func TestRandomReassign(t *testing.T) {
 		}
 	}
 }
+
+func TestBalancedAssign(t *testing.T) {
+	clusters := make([]*NodeInfo, 0, 3)
+	info1 := &datapb.DataNodeInfo{
+		Address:  "addr1",
+		Channels: make([]*datapb.ChannelStatus, 0, 10),
+	}
+	info2 := &datapb.DataNodeInfo{
+		Address:  "addr2",
+		Channels: make([]*datapb.ChannelStatus, 0, 10),
+	}
+	info3 := &datapb.DataNodeInfo{
+		Address:  "addr3",
+		Channels: make([]*datapb.ChannelStatus, 0, 10),
+	}
+
+	node1 := NewNodeInfo(context.TODO(), info1)
+	node2 := NewNodeInfo(context.TODO(), info2)
+	node3 := NewNodeInfo(context.TODO(), info3)
+	clusters = append(clusters, node1, node2, node3)
+
+	emptyClusters := make([]*NodeInfo, 0)
+
+	type testCase struct {
+		cluster      []*NodeInfo
+		channel      string
+		collectionID UniqueID
+		validator    func(c testCase, result []*NodeInfo) bool
+	}
+	hits := make(map[string]struct{})
+	for _, info := range clusters {
+		hits[info.Info.Address] = struct{}{}
+	}
+
+	strategyValidator := func(c testCase, result []*NodeInfo) bool {
+		for _, info := range result {
+			for _, channel := range info.Info.Channels {
+				if channel.CollectionID == c.collectionID && channel.Name == c.channel {
+					if _, ok := hits[info.Info.Address]; !ok {
+						return false
+					}
+					delete(hits, info.Info.Address)
+					return true
+				}
+			}
+		}
+		return false
+	}
+
+	cases := []testCase{
+		{
+			cluster:      emptyClusters,
+			channel:      "ch-01",
+			collectionID: 1,
+			validator: func(t testCase, result []*NodeInfo) bool {
+				return len(result) == 0
+			},
+		},
+		{
+			cluster:      clusters,
+			channel:      "ch-01",
+			collectionID: 1,
+			validator:    strategyValidator,
+		},
+		{
+			cluster:      clusters,
+			channel:      "ch-02",
+			collectionID: 1,
+			validator:    strategyValidator,
+		},
+		{
+			cluster:      clusters,
+			channel:      "ch-03",
+			collectionID: 1,
+			validator:    strategyValidator,
+		},
+	}
+	for _, c := range cases {
+		result := balancedAssignFunc(c.cluster, c.channel, c.collectionID)
+		if c.validator != nil {
+			assert.True(t, c.validator(c, result))
+		}
+		for _, info := range result {
+			for i, cluster := range clusters {
+				if cluster.Info.Address == info.Info.Address {
+					clusters[i] = info
+				}
+			}
+		}
+	}
+}

internal/datacoord/segment_allocation_policy.go

@@ -12,6 +12,7 @@
 package datacoord
 
 import (
+	"errors"
 	"sort"
 	"time"
 
@@ -24,10 +25,17 @@ import (
 type calUpperLimitPolicy func(schema *schemapb.CollectionSchema) (int, error)
 
 func calBySchemaPolicy(schema *schemapb.CollectionSchema) (int, error) {
+	if schema == nil {
+		return -1, errors.New("nil schema")
+	}
 	sizePerRecord, err := typeutil.EstimateSizePerRecord(schema)
 	if err != nil {
 		return -1, err
 	}
+	// check zero value, preventing panicking
+	if sizePerRecord == 0 {
+		return -1, errors.New("zero size record schema found")
+	}
 	threshold := Params.SegmentMaxSize * 1024 * 1024
 	return int(threshold / float64(sizePerRecord)), nil
 }

internal/datacoord/segment_allocation_policy_test.go

@@ -1,14 +1,142 @@
 package datacoord
 
 import (
+	"math/rand"
 	"testing"
 	"time"
 
+	"github.com/milvus-io/milvus/internal/proto/commonpb"
 	"github.com/milvus-io/milvus/internal/proto/datapb"
+	"github.com/milvus-io/milvus/internal/proto/schemapb"
 	"github.com/milvus-io/milvus/internal/util/tsoutil"
 	"github.com/stretchr/testify/assert"
 )
 
+func TestUpperLimitCalBySchema(t *testing.T) {
+	type testCase struct {
+		schema    *schemapb.CollectionSchema
+		expected  int
+		expectErr bool
+	}
+	testCases := []testCase{
+		{
+			schema:    nil,
+			expected:  -1,
+			expectErr: true,
+		},
+		{
+			schema: &schemapb.CollectionSchema{
+				Fields: []*schemapb.FieldSchema{},
+			},
+			expected:  -1,
+			expectErr: true,
+		},
+		{
+			schema: &schemapb.CollectionSchema{
+				Fields: []*schemapb.FieldSchema{
+					{
+						DataType: schemapb.DataType_FloatVector,
+						TypeParams: []*commonpb.KeyValuePair{
+							{Key: "dim", Value: "bad_dim"},
+						},
+					},
+				},
+			},
+			expected:  -1,
+			expectErr: true,
+		},
+		{
+			schema: &schemapb.CollectionSchema{
+				Fields: []*schemapb.FieldSchema{
+					{
+						DataType: schemapb.DataType_Int64,
+					},
+					{
+						DataType: schemapb.DataType_Int32,
+					},
+					{
+						DataType: schemapb.DataType_FloatVector,
+						TypeParams: []*commonpb.KeyValuePair{
+							{Key: "dim", Value: "128"},
+						},
+					},
+				},
+			},
+			expected:  int(Params.SegmentMaxSize * 1024 * 1024 / float64(524)),
+			expectErr: false,
+		},
+	}
+	for _, c := range testCases {
+		result, err := calBySchemaPolicy(c.schema)
+		if c.expectErr {
+			assert.NotNil(t, err)
+		} else {
+			assert.Equal(t, c.expected, result)
+		}
+	}
+}
+
+func TestGetChannelOpenSegCapacityPolicy(t *testing.T) {
+	p := getChannelOpenSegCapacityPolicy(3)
+	type testCase struct {
+		channel   string
+		segments  []*SegmentInfo
+		ts        Timestamp
+		validator func([]*SegmentInfo) bool
+	}
+	testCases := []testCase{
+		{
+			segments: []*SegmentInfo{},
+			ts:       tsoutil.ComposeTS(time.Now().Unix()/int64(time.Millisecond), rand.Int63n(1000)),
+			validator: func(result []*SegmentInfo) bool {
+				return len(result) == 0
+			},
+		},
+		{
+			segments: []*SegmentInfo{
+				{
+					SegmentInfo: &datapb.SegmentInfo{},
+				},
+				{
+					SegmentInfo: &datapb.SegmentInfo{},
+				},
+				{
+					SegmentInfo: &datapb.SegmentInfo{},
+				},
+				{
+					SegmentInfo: &datapb.SegmentInfo{},
+				},
+			},
+			ts: tsoutil.ComposeTS(time.Now().Unix()/int64(time.Millisecond), rand.Int63n(1000)),
+			validator: func(result []*SegmentInfo) bool {
+				return len(result) == 1
+			},
+		},
+	}
+	for _, c := range testCases {
+		result := p(c.channel, c.segments, c.ts)
+		if c.validator != nil {
+			assert.True(t, c.validator(result))
+		}
+	}
+}
+
+func TestSortSegmentsByLastExpires(t *testing.T) {
+	segs := make([]*SegmentInfo, 0, 10)
+	for i := 0; i < 10; i++ {
+		segs = append(segs,
+			&SegmentInfo{
+				SegmentInfo: &datapb.SegmentInfo{
+					LastExpireTime: uint64(rand.Int31n(100000)),
+				},
+			})
+	}
+	sortSegmentsByLastExpires(segs)
+	for i := 1; i < 10; i++ {
+		assert.True(t, segs[i-1].LastExpireTime <= segs[i].LastExpireTime)
+	}
+}
+
 func TestSealSegmentPolicy(t *testing.T) {
 	t.Run("test seal segment by lifetime", func(t *testing.T) {
 		lifetime := 2 * time.Second

internal/datacoord/server_test.go

@@ -12,6 +12,7 @@ package datacoord
 import (
 	"context"
 	"math/rand"
+	"os"
 	"path"
 	"strconv"
 	"testing"
@@ -28,8 +29,9 @@ import (
 	"go.etcd.io/etcd/clientv3"
 )
 
-func init() {
+func TestMain(m *testing.M) {
 	rand.Seed(time.Now().UnixNano())
+	os.Exit(m.Run())
 }
 
 func TestGetSegmentInfoChannel(t *testing.T) {

internal/datacoord/util.go

@@ -20,13 +20,16 @@ import (
 	"github.com/milvus-io/milvus/internal/proto/commonpb"
 )
 
+// Response response interface for verification
 type Response interface {
 	GetStatus() *commonpb.Status
 }
 
 var errNilResponse = errors.New("response is nil")
+var errNilStatusResponse = errors.New("response has nil status")
 var errUnknownResponseType = errors.New("unknown response type")
 
+// VerifyResponse verify grpc Response 1. check error is nil 2. check response.GetStatus() with status success
 func VerifyResponse(response interface{}, err error) error {
 	if err != nil {
 		return err
@@ -36,12 +39,19 @@ func VerifyResponse(response interface{}, err error) error {
 	}
 	switch resp := response.(type) {
 	case Response:
-		if resp.GetStatus().ErrorCode != commonpb.ErrorCode_Success {
-			return errors.New(resp.GetStatus().Reason)
+		// note that resp will not be nil here, since it's still a interface
+		if resp.GetStatus() == nil {
+			return errNilStatusResponse
+		}
+		if resp.GetStatus().GetErrorCode() != commonpb.ErrorCode_Success {
+			return errors.New(resp.GetStatus().GetReason())
 		}
 	case *commonpb.Status:
+		if resp == nil {
+			return errNilResponse
+		}
 		if resp.ErrorCode != commonpb.ErrorCode_Success {
-			return errors.New(resp.Reason)
+			return errors.New(resp.GetReason())
 		}
 	default:
 		return errUnknownResponseType
@@ -59,6 +69,7 @@ type LongTermChecker struct {
 	name string
 }
 
+// NewLongTermChecker create a long term checker specified name, checking interval and warning string to print
 func NewLongTermChecker(ctx context.Context, name string, d time.Duration, warn string) *LongTermChecker {
 	c := &LongTermChecker{
 		name: name,
@@ -69,6 +80,7 @@ func NewLongTermChecker(ctx context.Context, name string, d time.Duration, warn
 	return c
 }
 
+// Start starts the check process
 func (c *LongTermChecker) Start() {
 	c.t = time.NewTicker(c.d)
 	go func() {

internal/datacoord/util_test.go

@@ -0,0 +1,93 @@
+package datacoord
+
+import (
+	"errors"
+	"testing"
+
+	"github.com/milvus-io/milvus/internal/proto/commonpb"
+	"github.com/milvus-io/milvus/internal/proto/rootcoordpb"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestVerifyResponse(t *testing.T) {
+	type testCase struct {
+		resp       interface{}
+		err        error
+		expected   error
+		equalValue bool
+	}
+	cases := []testCase{
+		{
+			resp:       nil,
+			err:        errors.New("boom"),
+			expected:   errors.New("boom"),
+			equalValue: true,
+		},
+		{
+			resp:       nil,
+			err:        nil,
+			expected:   errNilResponse,
+			equalValue: false,
+		},
+		{
+			resp:       &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
+			err:        nil,
+			expected:   nil,
+			equalValue: false,
+		},
+		{
+			resp:       &commonpb.Status{ErrorCode: commonpb.ErrorCode_UnexpectedError, Reason: "r1"},
+			err:        nil,
+			expected:   errors.New("r1"),
+			equalValue: true,
+		},
+		{
+			resp:       (*commonpb.Status)(nil),
+			err:        nil,
+			expected:   errNilResponse,
+			equalValue: false,
+		},
+		{
+			resp: &rootcoordpb.AllocIDResponse{
+				Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
+			},
+			err:        nil,
+			expected:   nil,
+			equalValue: false,
+		},
+		{
+			resp: &rootcoordpb.AllocIDResponse{
+				Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_UnexpectedError, Reason: "r2"},
+			},
+			err:        nil,
+			expected:   errors.New("r2"),
+			equalValue: true,
+		},
+		{
+			resp:       &rootcoordpb.AllocIDResponse{},
+			err:        nil,
+			expected:   errNilStatusResponse,
+			equalValue: true,
+		},
+		{
+			resp:       (*rootcoordpb.AllocIDResponse)(nil),
+			err:        nil,
+			expected:   errNilStatusResponse,
+			equalValue: true,
+		},
+		{
+			resp:       struct{}{},
+			err:        nil,
+			expected:   errUnknownResponseType,
+			equalValue: false,
+		},
+	}
+	for _, c := range cases {
+		r := VerifyResponse(c.resp, c.err)
+		if c.equalValue {
+			assert.EqualValues(t, c.expected, r)
+		} else {
+			assert.Equal(t, c.expected, r)
+		}
+	}
+}