Add unittest for timestamp allocator (#7367)

Signed-off-by: dragondriver <jiquan.long@zilliz.com>

internal/proxy/mock.go

@@ -0,0 +1,39 @@
+// 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 proxy
+
+import (
+	"context"
+	"time"
+
+	"github.com/milvus-io/milvus/internal/proto/commonpb"
+
+	"github.com/milvus-io/milvus/internal/proto/rootcoordpb"
+)
+
+type mockTimestampAllocatorInterface struct {
+}
+
+func (tso *mockTimestampAllocatorInterface) AllocTimestamp(ctx context.Context, req *rootcoordpb.AllocTimestampRequest) (*rootcoordpb.AllocTimestampResponse, error) {
+	return &rootcoordpb.AllocTimestampResponse{
+		Status: &commonpb.Status{
+			ErrorCode: commonpb.ErrorCode_Success,
+			Reason:    "",
+		},
+		Timestamp: uint64(time.Now().UnixNano()),
+		Count:     req.Count,
+	}, nil
+}
+
+func newMockTimestampAllocatorInterface() timestampAllocatorInterface {
+	return &mockTimestampAllocatorInterface{}
+}

internal/proxy/timestamp.go

@@ -18,20 +18,24 @@ import (
 
 	"github.com/milvus-io/milvus/internal/proto/commonpb"
 	"github.com/milvus-io/milvus/internal/proto/rootcoordpb"
-	"github.com/milvus-io/milvus/internal/types"
 )
 
-type TimestampAllocator struct {
-	ctx       context.Context
-	rootCoord types.RootCoord
-	peerID    UniqueID
+// use timestampAllocatorInterface to keep TimestampAllocator testable
+type timestampAllocatorInterface interface {
+	AllocTimestamp(ctx context.Context, req *rootcoordpb.AllocTimestampRequest) (*rootcoordpb.AllocTimestampResponse, error)
 }
 
-func NewTimestampAllocator(ctx context.Context, rc types.RootCoord, peerID UniqueID) (*TimestampAllocator, error) {
+type TimestampAllocator struct {
+	ctx    context.Context
+	tso    timestampAllocatorInterface
+	peerID UniqueID
+}
+
+func NewTimestampAllocator(ctx context.Context, tso timestampAllocatorInterface, peerID UniqueID) (*TimestampAllocator, error) {
 	a := &TimestampAllocator{
-		ctx:       ctx,
-		peerID:    peerID,
-		rootCoord: rc,
+		ctx:    ctx,
+		peerID: peerID,
+		tso:    tso,
 	}
 	return a, nil
 }
@@ -48,7 +52,7 @@ func (ta *TimestampAllocator) Alloc(count uint32) ([]Timestamp, error) {
 		Count: count,
 	}
 
-	resp, err := ta.rootCoord.AllocTimestamp(ctx, req)
+	resp, err := ta.tso.AllocTimestamp(ctx, req)
 	defer cancel()
 
 	if err != nil {

internal/proxy/timestamp_test.go

@@ -0,0 +1,58 @@
+// 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 proxy
+
+import (
+	"context"
+	"math/rand"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNewTimestampAllocator(t *testing.T) {
+	ctx := context.Background()
+	tso := newMockTimestampAllocatorInterface()
+	peerID := UniqueID(getUniqueIntGeneratorIns().get())
+
+	tsAllocator, err := NewTimestampAllocator(ctx, tso, peerID)
+	assert.Nil(t, err)
+	assert.NotNil(t, tsAllocator)
+}
+
+func TestTimestampAllocator_Alloc(t *testing.T) {
+	ctx := context.Background()
+	tso := newMockTimestampAllocatorInterface()
+	peerID := UniqueID(getUniqueIntGeneratorIns().get())
+
+	tsAllocator, err := NewTimestampAllocator(ctx, tso, peerID)
+	assert.Nil(t, err)
+	assert.NotNil(t, tsAllocator)
+
+	count := rand.Uint32()%100 + 1
+	ret, err := tsAllocator.Alloc(count)
+	assert.Nil(t, err)
+	assert.Equal(t, int(count), len(ret))
+}
+
+func TestTimestampAllocator_AllocOne(t *testing.T) {
+	ctx := context.Background()
+	tso := newMockTimestampAllocatorInterface()
+	peerID := UniqueID(getUniqueIntGeneratorIns().get())
+
+	tsAllocator, err := NewTimestampAllocator(ctx, tso, peerID)
+	assert.Nil(t, err)
+	assert.NotNil(t, tsAllocator)
+
+	_, err = tsAllocator.AllocOne()
+	assert.Nil(t, err)
+}