Add unit test

Signed-off-by: Filip Petkovski <filip.petkovsky@gmail.com>

pkg/query/iter_test.go

@@ -0,0 +1,74 @@
+// Copyright (c) The Thanos Authors.
+// Licensed under the Apache License 2.0.
+
+package query
+
+import (
+	"testing"
+
+	"github.com/efficientgo/core/testutil"
+	"github.com/prometheus/prometheus/model/histogram"
+	"github.com/prometheus/prometheus/tsdb/tsdbutil"
+)
+
+func TestNativeHistogramDedup(t *testing.T) {
+	cases := []struct {
+		samples []sample
+		tcase   string
+	}{
+		{
+			tcase: "unknown counter reset chunk",
+			samples: []sample{
+				{t: 10000, h: makeHistWithHint(1, histogram.UnknownCounterReset)},
+				{t: 20000, h: makeHistWithHint(2, histogram.UnknownCounterReset)},
+			},
+		},
+		{
+			tcase: "not counter reset chunk",
+			samples: []sample{
+				{t: 10000, h: makeHistWithHint(1, histogram.NotCounterReset)},
+				{t: 20000, h: makeHistWithHint(2, histogram.NotCounterReset)},
+			},
+		},
+		{
+			tcase: "counter reset chunk",
+			samples: []sample{
+				{t: 10000, h: makeHistWithHint(1, histogram.CounterReset)},
+				{t: 20000, h: makeHistWithHint(2, histogram.NotCounterReset)},
+			},
+		},
+	}
+
+	for _, c := range cases {
+		t.Run(c.tcase, func(t *testing.T) {
+			// When the first sample is read, the counter reset hint for the second sample
+			// does not need to be reset.
+			t.Run("read_first_sample", func(t *testing.T) {
+				it := newHistogramResetDetector(newMockedSeriesIterator(c.samples))
+				it.Next()
+				_, h := it.AtHistogram()
+				testutil.Equals(t, c.samples[0].h.CounterResetHint, h.CounterResetHint)
+
+				it.Next()
+				_, h = it.AtHistogram()
+				testutil.Equals(t, c.samples[1].h.CounterResetHint, h.CounterResetHint)
+			})
+
+			// When the first sample is not read, the counter reset hint for the second
+			// sample should be reset.
+			t.Run("skip_first_sample", func(t *testing.T) {
+				it := newHistogramResetDetector(newMockedSeriesIterator(c.samples))
+				it.Next()
+				it.Next()
+				_, h := it.AtHistogram()
+				testutil.Equals(t, h.CounterResetHint, histogram.UnknownCounterReset)
+			})
+		})
+	}
+}
+
+func makeHistWithHint(i int, hint histogram.CounterResetHint) *histogram.Histogram {
+	h := tsdbutil.GenerateTestHistogram(i)
+	h.CounterResetHint = hint
+	return h
+}

pkg/query/querier_test.go

@@ -1020,11 +1020,13 @@ func (s *mockedSeriesIterator) Seek(t int64) chunkenc.ValueType {
 		return s.samples[n].t >= t
 	})
 
-	if s.cur < len(s.samples) {
-		return chunkenc.ValFloat
+	if s.cur >= len(s.samples) {
+		return chunkenc.ValNone
 	}
-
-	return chunkenc.ValNone
+	if s.samples[s.cur].h != nil {
+		return chunkenc.ValHistogram
+	}
+	return chunkenc.ValFloat
 }
 
 func (s *mockedSeriesIterator) At() (t int64, v float64) {
@@ -1034,7 +1036,8 @@ func (s *mockedSeriesIterator) At() (t int64, v float64) {
 
 // TODO(rabenhorst): Needs to be implemented for native histogram support.
 func (s *mockedSeriesIterator) AtHistogram(*histogram.Histogram) (int64, *histogram.Histogram) {
-	panic("not implemented")
+	sample := s.samples[s.cur]
+	return sample.t, sample.h
 }
 
 func (s *mockedSeriesIterator) AtFloatHistogram(*histogram.FloatHistogram) (int64, *histogram.FloatHistogram) {
@@ -1047,11 +1050,13 @@ func (s *mockedSeriesIterator) AtT() int64 {
 
 func (s *mockedSeriesIterator) Next() chunkenc.ValueType {
 	s.cur++
-	if s.cur < len(s.samples) {
-		return chunkenc.ValFloat
+	if s.cur >= len(s.samples) {
+		return chunkenc.ValNone
 	}
-
-	return chunkenc.ValNone
+	if s.samples[s.cur].h != nil {
+		return chunkenc.ValHistogram
+	}
+	return chunkenc.ValFloat
 }
 
 func (s *mockedSeriesIterator) Err() error { return nil }
@@ -1205,14 +1210,20 @@ func TestQuerierWithDedupUnderstoodByPromQL_Rate(t *testing.T) {
 const hackyStaleMarker = float64(-99999999)
 
 func expandSeries(t testing.TB, it chunkenc.Iterator) (res []sample) {
-	for it.Next() != chunkenc.ValNone {
-		t, v := it.At()
-		// Nan != Nan, so substitute for another value.
-		// This is required for testutil.Equals to work deterministically.
-		if math.IsNaN(v) {
-			v = hackyStaleMarker
+	for valType := it.Next(); valType != chunkenc.ValNone; valType = it.Next() {
+		switch valType {
+		case chunkenc.ValFloat:
+			t, v := it.At()
+			// Nan != Nan, so substitute for another value.
+			// This is required for testutil.Equals to work deterministically.
+			if math.IsNaN(v) {
+				v = hackyStaleMarker
+			}
+			res = append(res, sample{t: t, v: v})
+		case chunkenc.ValHistogram:
+			t, h := it.AtHistogram(nil)
+			res = append(res, sample{t: t, h: h})
 		}
-		res = append(res, sample{t: t, v: v})
 	}
 	testutil.Ok(t, it.Err())
 	return res