From 1b48e84c35d8dcd510e19fc87d37e8a7f815020e Mon Sep 17 00:00:00 2001
From: XuanYang-cn <xuan.yang@zilliz.com>
Date: Fri, 13 Nov 2020 15:44:02 +0800
Subject: [PATCH] Refactor master scheduler: has_collection

Signed-off-by: XuanYang-cn <xuan.yang@zilliz.com>
---
 internal/master/collection_task_test.go | 47 ++++++++++++++++++++++++-
 internal/master/grpc_service.go         |  3 +-
 2 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/internal/master/collection_task_test.go b/internal/master/collection_task_test.go
index 9047e4fc4c..d81a785517 100644
--- a/internal/master/collection_task_test.go
+++ b/internal/master/collection_task_test.go
@@ -113,6 +113,37 @@ func TestMaster_CreateCollectionTask(t *testing.T) {
 	assert.Nil(t, err)
 	assert.Equal(t, st.ErrorCode, commonpb.ErrorCode_SUCCESS)
 
+	reqHasCollection := internalpb.HasCollectionRequest{
+		MsgType:   internalpb.MsgType_kHasCollection,
+		ReqID:     1,
+		Timestamp: 11,
+		ProxyID:   1,
+		CollectionName: &servicepb.CollectionName{
+			CollectionName: "col1",
+		},
+	}
+
+	// HasCollection "col1" is true
+	log.Printf("... [Has] collection col1\n")
+	boolResp, err := cli.HasCollection(ctx, &reqHasCollection)
+	assert.Nil(t, err)
+	assert.Equal(t, true, boolResp.Value)
+	assert.Equal(t, boolResp.Status.ErrorCode, commonpb.ErrorCode_SUCCESS)
+
+	// HasCollection "colNotExist" is false
+	reqHasCollection.CollectionName.CollectionName = "colNotExist"
+	boolResp, err = cli.HasCollection(ctx, &reqHasCollection)
+	assert.Nil(t, err)
+	assert.Equal(t, boolResp.Value, false)
+	assert.Equal(t, boolResp.Status.ErrorCode, commonpb.ErrorCode_SUCCESS)
+
+	// HasCollection error
+	reqHasCollection.Timestamp = Timestamp(10)
+	reqHasCollection.CollectionName.CollectionName = "col1"
+	boolResp, err = cli.HasCollection(ctx, &reqHasCollection)
+	assert.Nil(t, err)
+	assert.NotEqual(t, boolResp.Status.ErrorCode, commonpb.ErrorCode_SUCCESS)
+
 	collMeta, err := svr.mt.GetCollectionByName(sch.Name)
 	assert.Nil(t, err)
 	t.Logf("collection id = %d", collMeta.ID)
@@ -150,8 +181,8 @@ func TestMaster_CreateCollectionTask(t *testing.T) {
 	st, err = cli.CreateCollection(ctx, &req)
 	assert.Nil(t, err)
 	assert.NotEqual(t, st.ErrorCode, commonpb.ErrorCode_SUCCESS)
-	// ------------------------------DropCollectionTask---------------------------
 
+	// ------------------------------DropCollectionTask---------------------------
 	log.Printf("... [Drop] collection col1\n")
 	ser := servicepb.CollectionName{CollectionName: "col1"}
 
@@ -163,6 +194,7 @@ func TestMaster_CreateCollectionTask(t *testing.T) {
 		CollectionName: &ser,
 	}
 
+	// DropCollection
 	st, err = cli.DropCollection(ctx, &reqDrop)
 	assert.Nil(t, err)
 	assert.Equal(t, st.ErrorCode, commonpb.ErrorCode_SUCCESS)
@@ -170,6 +202,14 @@ func TestMaster_CreateCollectionTask(t *testing.T) {
 	collMeta, err = svr.mt.GetCollectionByName(sch.Name)
 	assert.NotNil(t, err)
 
+	// HasCollection "col1" is false
+	reqHasCollection.Timestamp = Timestamp(11)
+	reqHasCollection.CollectionName.CollectionName = "col1"
+	boolResp, err = cli.HasCollection(ctx, &reqHasCollection)
+	assert.Nil(t, err)
+	assert.Equal(t, false, boolResp.Value)
+	assert.Equal(t, commonpb.ErrorCode_SUCCESS, boolResp.Status.ErrorCode)
+
 	// Drop again
 	st, err = cli.DropCollection(ctx, &reqDrop)
 	assert.Nil(t, err)
@@ -182,5 +222,10 @@ func TestMaster_CreateCollectionTask(t *testing.T) {
 	log.Printf(st.Reason)
 	assert.Equal(t, st.ErrorCode, commonpb.ErrorCode_SUCCESS)
 
+	boolResp, err = cli.HasCollection(ctx, &reqHasCollection)
+	assert.Nil(t, err)
+	assert.Equal(t, true, boolResp.Value)
+	assert.Equal(t, commonpb.ErrorCode_SUCCESS, boolResp.Status.ErrorCode)
+
 	svr.Close()
 }
diff --git a/internal/master/grpc_service.go b/internal/master/grpc_service.go
index 3048d403ae..b399a91883 100644
--- a/internal/master/grpc_service.go
+++ b/internal/master/grpc_service.go
@@ -93,7 +93,7 @@ func (s *Master) HasCollection(ctx context.Context, in *internalpb.HasCollection
 
 	response := &servicepb.BoolResponse{
 		Status: st,
-		Value:  t.(*hasCollectionTask).hasCollection,
+		Value:  false,
 	}
 
 	var err = s.scheduler.Enqueue(t)
@@ -109,6 +109,7 @@ func (s *Master) HasCollection(ctx context.Context, in *internalpb.HasCollection
 	}
 
 	st.ErrorCode = commonpb.ErrorCode_SUCCESS
+	response.Value = t.(*hasCollectionTask).hasCollection
 	return response, nil
 }