diff --git a/pkg/querier/queryrange/marshal.go b/pkg/querier/queryrange/marshal.go index 0e72b15140e38..ef782d977b206 100644 --- a/pkg/querier/queryrange/marshal.go +++ b/pkg/querier/queryrange/marshal.go @@ -234,6 +234,8 @@ func QueryResponseWrap(res queryrangebase.Response) (*QueryResponse, error) { p.Response = &QueryResponse_Labels{response} case *IndexStatsResponse: p.Response = &QueryResponse_Stats{response} + case *VolumeResponse: + p.Response = &QueryResponse_Volume{response} case *TopKSketchesResponse: p.Response = &QueryResponse_TopkSketches{response} case *QuantileSketchResponse: diff --git a/pkg/querier/queryrange/marshal_test.go b/pkg/querier/queryrange/marshal_test.go index 569a1af1a4b98..e83ea607af877 100644 --- a/pkg/querier/queryrange/marshal_test.go +++ b/pkg/querier/queryrange/marshal_test.go @@ -43,3 +43,26 @@ func TestResultToResponse(t *testing.T) { }) } } + +func TestResponseWrap(t *testing.T) { + for _, tt := range []struct { + name string + response queryrangebase.Response + expected isQueryResponse_Response + }{ + {"volume", &VolumeResponse{}, &QueryResponse_Volume{}}, + {"series", &LokiSeriesResponse{}, &QueryResponse_Series{}}, + {"label", &LokiLabelNamesResponse{}, &QueryResponse_Labels{}}, + {"stats", &IndexStatsResponse{}, &QueryResponse_Stats{}}, + {"prom", &LokiPromResponse{}, &QueryResponse_Prom{}}, + {"streams", &LokiResponse{}, &QueryResponse_Streams{}}, + {"topk", &TopKSketchesResponse{}, &QueryResponse_TopkSketches{}}, + {"quantile", &QuantileSketchResponse{}, &QueryResponse_QuantileSketches{}}, + } { + t.Run(tt.name, func(t *testing.T) { + actual, err := QueryResponseWrap(tt.response) + require.NoError(t, err) + require.IsType(t, tt.expected, actual.Response) + }) + } +}