diff --git a/internal/msgstream/msg.go b/internal/msgstream/msg.go index d5fdaa2930..69b31da9ba 100644 --- a/internal/msgstream/msg.go +++ b/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 diff --git a/internal/msgstream/msg_test.go b/internal/msgstream/msg_test.go index 1ea0147984..9c8092406f 100644 --- a/internal/msgstream/msg_test.go +++ b/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(),