Add Unittest for allocator (#5781)

Signed-off-by: yangxuan <xuan.yang@zilliz.com>

configs/advanced/data_node.yaml

@@ -19,4 +19,3 @@ dataNode:
   flush:
     # max buffer size to flush
     insertBufSize: 32000 # number of rows
-    ddBufSize: 20

configs/milvus.yaml

@@ -19,8 +19,6 @@ etcd:
   rootPath: by-dev
   metaSubPath: meta # metaRootPath = rootPath + '/' + metaSubPath
   kvSubPath: kv # kvRootPath = rootPath + '/' + kvSubPath
-  segFlushMetaSubPath: datanode/segment # Full Path = rootPath/metaSubPath/segFlushMetaSubPath
-  ddlFlushMetaSubPath: datanode/ddl     # Full Path = rootPath/metaSubPath/ddlFlushMetaSubPath
   segmentBinlogSubPath: dataservice/binlog/segment  # Full Path = rootPath/metaSubPath/segmentBinlogSubPath
   collectionBinlogSubPath: dataservice/binlog/collection # Full Path = rootPath/metaSubPath/collectionBinglogSubPath
   flushStreamPosSubPath: dataservice/flushstream # Full path = rootPath/metaSubPath/flushStreamPosSubPath

internal/datanode/allocator.go

@@ -13,6 +13,7 @@ package datanode
 
 import (
 	"context"
+	"errors"
 	"path"
 	"strconv"
 
@@ -50,9 +51,15 @@ func (alloc *allocator) allocID() (UniqueID, error) {
 		},
 		Count: 1,
 	})
+
+	if resp.Status.ErrorCode != commonpb.ErrorCode_Success {
+		return 0, errors.New(resp.Status.GetReason())
+	}
+
 	if err != nil {
 		return 0, err
 	}
+
 	return resp.ID, nil
 }
 

internal/datanode/allocator_test.go

@@ -22,7 +22,83 @@ func TestAllocator_Basic(t *testing.T) {
 	allocator := newAllocator(ms)
 
 	t.Run("Test allocID", func(t *testing.T) {
+		ms.setID(666)
 		_, err := allocator.allocID()
 		assert.NoError(t, err)
 	})
+
+	t.Run("Test genKey", func(t *testing.T) {
+		ms.setID(666)
+
+		type in struct {
+			isalloc bool
+			ids     []UniqueID
+		}
+
+		type out struct {
+			key string
+			err error
+		}
+
+		type Test struct {
+			in
+			out
+		}
+
+		tests := []Test{
+			{in{true, []UniqueID{}}, out{"666", nil}},
+			{in{true, []UniqueID{1}}, out{"1/666", nil}},
+			{in{true, make([]UniqueID, 0)}, out{"666", nil}},
+			{in{false, []UniqueID{}}, out{"", nil}},
+			{in{false, []UniqueID{1, 2, 3}}, out{"1/2/3", nil}},
+			{in{false, []UniqueID{1}}, out{"1", nil}},
+			{in{false, []UniqueID{2, 2, 2}}, out{"2/2/2", nil}},
+		}
+
+		for i, test := range tests {
+			key, err := allocator.genKey(test.in.isalloc, test.in.ids...)
+
+			assert.Equalf(t, test.out.key, key, "#%d", i)
+			assert.Equalf(t, test.out.err, err, "#%d", i)
+		}
+
+		// Status.ErrorCode != Success
+		ms.setID(0)
+		tests = []Test{
+			{in{true, []UniqueID{}}, out{}},
+			{in{true, []UniqueID{1}}, out{}},
+			{in{true, make([]UniqueID, 0)}, out{}},
+		}
+
+		for i, test := range tests {
+			_, err := allocator.genKey(test.in.isalloc, test.in.ids...)
+			assert.Errorf(t, err, "number: %d", i)
+		}
+
+		// Grpc error
+		ms.setID(-1)
+		tests = []Test{
+			{in{true, make([]UniqueID, 0)}, out{}},
+			{in{true, []UniqueID{1}}, out{}},
+			{in{true, make([]UniqueID, 0)}, out{}},
+		}
+
+		for i, test := range tests {
+			_, err := allocator.genKey(test.in.isalloc, test.in.ids...)
+			assert.Errorf(t, err, "number: %d", i)
+		}
+
+		// MasterService's unavailability doesn't affects genKey when alloc == false
+		tests = []Test{
+			{in{false, []UniqueID{1, 2, 3}}, out{"1/2/3", nil}},
+			{in{false, []UniqueID{1}}, out{"1", nil}},
+			{in{false, []UniqueID{2, 2, 2}}, out{"2/2/2", nil}},
+		}
+
+		for i, test := range tests {
+			key, err := allocator.genKey(test.in.isalloc, test.in.ids...)
+			assert.Equalf(t, test.out.key, key, "#%d", i)
+			assert.Equalf(t, test.out.err, err, "#%d", i)
+		}
+	})
 }

internal/datanode/collection_replica_test.go

@@ -16,8 +16,25 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
+	"github.com/milvus-io/milvus/internal/proto/internalpb"
+	"github.com/milvus-io/milvus/internal/types"
 )
 
+func newCollectionSegmentReplica(ms types.MasterService, collectionID UniqueID) *CollectionSegmentReplica {
+	metaService := newMetaService(ms, collectionID)
+	segments := make(map[UniqueID]*Segment)
+
+	replica := &CollectionSegmentReplica{
+		segments:       segments,
+		collection:     &Collection{id: collectionID},
+		metaService:    metaService,
+		startPositions: make(map[UniqueID][]*internalpb.MsgPosition),
+		endPositions:   make(map[UniqueID][]*internalpb.MsgPosition),
+	}
+	return replica
+}
+
 func TestReplica_Collection(t *testing.T) {
 	collID := UniqueID(100)
 

internal/datanode/mock_test.go

@@ -15,6 +15,7 @@ import (
 	"bytes"
 	"context"
 	"encoding/binary"
+	"errors"
 	"math"
 	"math/rand"
 	"path"
@@ -484,6 +485,8 @@ func (alloc *AllocatorFactory) genKey(isalloc bool, ids ...UniqueID) (key string
 	return
 }
 
+// If id == 0, AllocID will return not successful status
+// If id == -1, AllocID will return err
 func (m *MasterServiceFactory) setID(id UniqueID) {
 	m.ID = id // GOOSE TODO: random ID generator
 }
@@ -498,9 +501,22 @@ func (m *MasterServiceFactory) setCollectionName(name string) {
 
 func (m *MasterServiceFactory) AllocID(ctx context.Context, in *masterpb.AllocIDRequest) (*masterpb.AllocIDResponse, error) {
 	resp := &masterpb.AllocIDResponse{
-		Status: &commonpb.Status{},
-		ID:     m.ID,
+		Status: &commonpb.Status{
+			ErrorCode: commonpb.ErrorCode_UnexpectedError,
+		}}
+
+	if m.ID == 0 {
+		resp.Status.Reason = "Zero ID"
+		return resp, nil
 	}
+
+	if m.ID == -1 {
+		resp.Status.ErrorCode = commonpb.ErrorCode_Success
+		return resp, errors.New(resp.Status.GetReason())
+	}
+
+	resp.ID = m.ID
+	resp.Status.ErrorCode = commonpb.ErrorCode_Success
 	return resp, nil
 }