From a2f2284932c6d0d5490cdfff4a26deeb84660fa8 Mon Sep 17 00:00:00 2001
From: "cai.zhang" <cai.zhang@zilliz.com>
Date: Tue, 10 Nov 2020 15:56:57 +0800
Subject: [PATCH] Add code coverage for go files

Signed-off-by: cai.zhang <cai.zhang@zilliz.com>
---
 internal/msgstream/msgstream.go | 9 ++++++++-
 scripts/run_go_unittest.sh      | 2 +-
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/internal/msgstream/msgstream.go b/internal/msgstream/msgstream.go
index 3754b2dec6..9af3795fc2 100644
--- a/internal/msgstream/msgstream.go
+++ b/internal/msgstream/msgstream.go
@@ -31,6 +31,7 @@ type MsgStream interface {
 	Produce(*MsgPack) error
 	Broadcast(*MsgPack) error
 	Consume() *MsgPack
+	Chan() <-chan *MsgPack
 }
 
 type PulsarMsgStream struct {
@@ -108,7 +109,9 @@ func (ms *PulsarMsgStream) Start() {
 
 func (ms *PulsarMsgStream) Close() {
 	ms.streamCancel()
-	ms.wait.Wait()
+	if ms.wait != nil {
+		ms.wait.Wait()
+	}
 
 	for _, producer := range ms.producers {
 		if producer != nil {
@@ -252,6 +255,10 @@ func (ms *PulsarMsgStream) bufMsgPackToChannel() {
 	}
 }
 
+func (ms *PulsarMsgStream) Chan() <-chan *MsgPack {
+	return ms.receiveBuf
+}
+
 type PulsarTtMsgStream struct {
 	PulsarMsgStream
 	inputBuf      []*TsMsg
diff --git a/scripts/run_go_unittest.sh b/scripts/run_go_unittest.sh
index 0f5e71db1d..765f082684 100755
--- a/scripts/run_go_unittest.sh
+++ b/scripts/run_go_unittest.sh
@@ -13,4 +13,4 @@ SCRIPTS_DIR="$( cd -P "$( dirname "$SOURCE" )" && pwd )"
 # ignore Minio,S3 unittes
 MILVUS_DIR="${SCRIPTS_DIR}/../internal/"
 echo $MILVUS_DIR
-go test "${MILVUS_DIR}/kv/..." "${MILVUS_DIR}/msgstream/..." "${MILVUS_DIR}/master/..." "${MILVUS_DIR}/reader/..." -failfast
+go test -cover "${MILVUS_DIR}/kv/..." "${MILVUS_DIR}/msgstream/..." "${MILVUS_DIR}/master/..." "${MILVUS_DIR}/reader/..." -failfast