From 845359d4ab876ccad1e3dc9c1efcd51850dbed86 Mon Sep 17 00:00:00 2001
From: Christian Haudum <christian.haudum@gmail.com>
Date: Tue, 9 Jul 2024 13:43:47 +0200
Subject: [PATCH] test: Add range aggregation test for instant queries (#13447)

This PR adds test coverage for the `streamRangeVectorIterator`, which is used by instant queries to evaluate the single step of the query.

Signed-off-by: Christian Haudum <christian.haudum@gmail.com>
Co-authored-by: Ashwanth <iamashwanth@gmail.com>
---
 pkg/logql/range_vector_test.go | 136 ++++++++++++++++++++++++++++++++-
 1 file changed, 134 insertions(+), 2 deletions(-)

diff --git a/pkg/logql/range_vector_test.go b/pkg/logql/range_vector_test.go
index fb28ea5c9c..0fe81e7e2a 100644
--- a/pkg/logql/range_vector_test.go
+++ b/pkg/logql/range_vector_test.go
@@ -59,11 +59,11 @@ func newfakePeekingSampleIterator(samples []logproto.Sample) iter.PeekingSampleI
 }
 
 func newSample(t time.Time, v float64, metric labels.Labels) promql.Sample {
-	return promql.Sample{Metric: metric, T: t.UnixNano() / 1e+6, F: v}
+	return promql.Sample{Metric: metric, T: t.UnixMilli(), F: v}
 }
 
 func newPoint(t time.Time, v float64) promql.FPoint {
-	return promql.FPoint{T: t.UnixNano() / 1e+6, F: v}
+	return promql.FPoint{T: t.UnixMilli(), F: v}
 }
 
 func Benchmark_RangeVectorIteratorCompare(b *testing.B) {
@@ -217,6 +217,138 @@ func Benchmark_RangeVectorIterator(b *testing.B) {
 
 }
 
+func Test_RangeVectorIterator_InstantQuery(t *testing.T) {
+	now := time.Date(2024, 7, 3, 0, 0, 0, 0, time.UTC)
+
+	samples := []logproto.Sample{
+		{Timestamp: now.Add(-2 * time.Hour).UnixNano(), Value: 1.23},
+		{Timestamp: now.Add(-1 * time.Hour).UnixNano(), Value: 2.34},
+		{Timestamp: now.UnixNano(), Value: 3.45},
+	}
+
+	tests := []struct {
+		selRange       time.Duration
+		now            time.Time
+		expectedVector promql.Vector
+		expectedTs     time.Time
+	}{
+		{
+			// query range:    (----]
+			// samples:        x        x        x
+			selRange:       30 * time.Minute,
+			now:            now.Add(-90 * time.Minute),
+			expectedVector: []promql.Sample{},
+			expectedTs:     now.Add(-90 * time.Minute),
+		},
+		{
+			// query range:    (--------]
+			// samples:        x        x        x
+			selRange: 60 * time.Minute,
+			now:      now.Add(-60 * time.Minute),
+			expectedVector: []promql.Sample{
+				newSample(now.Add(-60*time.Minute), 1, labelFoo),
+				newSample(now.Add(-60*time.Minute), 1, labelBar),
+			},
+			expectedTs: now.Add(-60 * time.Minute),
+		},
+		{
+			// query range:        (----]
+			// samples:        x        x        x
+			selRange: 30 * time.Minute,
+			now:      now.Add(-60 * time.Minute),
+			expectedVector: []promql.Sample{
+				newSample(now.Add(-60*time.Minute), 1, labelFoo),
+				newSample(now.Add(-60*time.Minute), 1, labelBar),
+			},
+			expectedTs: now.Add(-60 * time.Minute),
+		},
+		{
+			// query range:        (---------]
+			// samples:        x        x        x
+			selRange: 60 * time.Minute,
+			now:      now.Add(-30 * time.Minute),
+			expectedVector: []promql.Sample{
+				newSample(now.Add(-30*time.Minute), 1, labelFoo),
+				newSample(now.Add(-30*time.Minute), 1, labelBar),
+			},
+			expectedTs: now.Add(-30 * time.Minute),
+		},
+		{
+			// query range:             (----]
+			// samples:        x        x        x
+			selRange:       30 * time.Minute,
+			now:            now.Add(-30 * time.Minute),
+			expectedVector: []promql.Sample{},
+			expectedTs:     now.Add(-30 * time.Minute),
+		},
+		{
+			// query range:             (--------]
+			// samples:        x        x        x
+			selRange: 60 * time.Minute,
+			now:      now,
+			expectedVector: []promql.Sample{
+				newSample(now, 1, labelFoo),
+				newSample(now, 1, labelBar),
+			},
+			expectedTs: now,
+		},
+		{
+			// query range:                 (----]
+			// samples:        x        x        x
+			selRange: 30 * time.Minute,
+			now:      now,
+			expectedVector: []promql.Sample{
+				newSample(now, 1, labelFoo),
+				newSample(now, 1, labelBar),
+			},
+			expectedTs: now,
+		},
+		{
+			// query range:    (-----------------]
+			// samples:        x        x        x
+			selRange: 120 * time.Minute,
+			now:      now,
+			expectedVector: []promql.Sample{
+				newSample(now, 2, labelFoo),
+				newSample(now, 2, labelBar),
+			},
+			expectedTs: now,
+		},
+		{
+			// query range:               (----]
+			// samples:        x        x        x
+			selRange:       30 * time.Minute,
+			now:            now.Add(-15 * time.Minute),
+			expectedVector: []promql.Sample{},
+			expectedTs:     now.Add(-15 * time.Minute),
+		},
+	}
+
+	for i, tt := range tests {
+		t.Run(
+			fmt.Sprintf("%d logs[%s] - start %s - end %s", i, tt.selRange, tt.now.Add(-tt.selRange), tt.now),
+			func(t *testing.T) {
+				it, err := newRangeVectorIterator(
+					newfakePeekingSampleIterator(samples),
+					&syntax.RangeAggregationExpr{Operation: syntax.OpRangeTypeCount},
+					tt.selRange.Nanoseconds(),
+					0,                 // step
+					tt.now.UnixNano(), // start
+					tt.now.UnixNano(), // end
+					0,                 // offset
+				)
+				require.NoError(t, err)
+
+				hasNext := it.Next()
+				require.True(t, hasNext)
+
+				ts, v := it.At()
+				require.ElementsMatch(t, tt.expectedVector, v)
+				require.Equal(t, tt.expectedTs.UnixMilli(), ts)
+			})
+	}
+}
+
 func Test_RangeVectorIterator(t *testing.T) {
 	tests := []struct {
 		selRange        int64