Add unit tests for msg.go (#8227)

Signed-off-by: Xiangyu Wang <xiangyu.wang@zilliz.com>

internal/msgstream/msg.go

@@ -145,7 +145,6 @@ func (it *InsertMsg) Unmarshal(input MarshalType) (TsMsg, error) {
 }
 
 /////////////////////////////////////////Delete//////////////////////////////////////////
-// TODO(wxyu): comment it until really needed
 type DeleteMsg struct {
 	BaseMsg
 	internalpb.DeleteRequest
@@ -763,7 +762,7 @@ func (dp *DropPartitionMsg) Unmarshal(input MarshalType) (TsMsg, error) {
 }
 
 /////////////////////////////////////////LoadIndex//////////////////////////////////////////
-// TODO(wxyu): comment it until really needed
+// FIXME(wxyu): comment it until really needed
 /*
 type LoadIndexMsg struct {
 	BaseMsg

internal/msgstream/msg_test.go

@@ -454,6 +454,51 @@ func TestTimeTickMsg_Unmarshal_IllegalParameter(t *testing.T) {
 	assert.Nil(t, tsMsg)
 }
 
+func TestQueryNodeStatsMsg(t *testing.T) {
+	queryNodeStatsMsg := &QueryNodeStatsMsg{
+		BaseMsg: generateBaseMsg(),
+		QueryNodeStats: internalpb.QueryNodeStats{
+			Base: &commonpb.MsgBase{
+				MsgType:   commonpb.MsgType_TimeTick,
+				MsgID:     1,
+				Timestamp: 2,
+				SourceID:  3,
+			},
+			SegStats:   []*internalpb.SegmentStats{},
+			FieldStats: []*internalpb.FieldStats{},
+		},
+	}
+
+	assert.NotNil(t, queryNodeStatsMsg.TraceCtx())
+
+	ctx := context.Background()
+	queryNodeStatsMsg.SetTraceCtx(ctx)
+	assert.Equal(t, ctx, queryNodeStatsMsg.TraceCtx())
+
+	assert.Equal(t, int64(1), queryNodeStatsMsg.ID())
+	assert.Equal(t, commonpb.MsgType_TimeTick, queryNodeStatsMsg.Type())
+	assert.Equal(t, int64(3), queryNodeStatsMsg.SourceID())
+
+	bytes, err := queryNodeStatsMsg.Marshal(queryNodeStatsMsg)
+	assert.Nil(t, err)
+
+	tsMsg, err := queryNodeStatsMsg.Unmarshal(bytes)
+	assert.Nil(t, err)
+
+	queryNodeStatsMsg2, ok := tsMsg.(*QueryNodeStatsMsg)
+	assert.True(t, ok)
+	assert.Equal(t, int64(1), queryNodeStatsMsg2.ID())
+	assert.Equal(t, commonpb.MsgType_TimeTick, queryNodeStatsMsg2.Type())
+	assert.Equal(t, int64(3), queryNodeStatsMsg2.SourceID())
+}
+
+func TestQueryNodeStatsMsg_Unmarshal_IllegalParameter(t *testing.T) {
+	queryNodeStatsMsg := &QueryNodeStatsMsg{}
+	tsMsg, err := queryNodeStatsMsg.Unmarshal(10)
+	assert.NotNil(t, err)
+	assert.Nil(t, tsMsg)
+}
+
 func TestSegmentStatisticsMsg(t *testing.T) {
 	segmentStatisticsMsg := &SegmentStatisticsMsg{
 		BaseMsg: generateBaseMsg(),