From 442f24e09927957936dd2a5b2c5c0817da0674cf Mon Sep 17 00:00:00 2001
From: Arthur Silva Sens <arthursens2005@gmail.com>
Date: Mon, 2 Sep 2024 17:30:37 -0300
Subject: [PATCH] chore: Simplify TestHeadAppender_AppendCTZeroSample (#14812)

Signed-off-by: Arthur Silva Sens <arthursens2005@gmail.com>
---
 tsdb/head_test.go | 81 +++++++++++++++++++++--------------------------
 1 file changed, 36 insertions(+), 45 deletions(-)

diff --git a/tsdb/head_test.go b/tsdb/head_test.go
index b32b37fca..18ec4f0ac 100644
--- a/tsdb/head_test.go
+++ b/tsdb/head_test.go
@@ -33,7 +33,6 @@ import (
 	"github.com/google/go-cmp/cmp"
 	"github.com/prometheus/client_golang/prometheus"
 	prom_testutil "github.com/prometheus/client_golang/prometheus/testutil"
-	"github.com/prometheus/common/model"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/atomic"
 	"golang.org/x/sync/errgroup"
@@ -5945,16 +5944,16 @@ func TestHeadAppender_AppendCTZeroSample(t *testing.T) {
 	for _, tc := range []struct {
 		name              string
 		appendableSamples []appendableSamples
-		expectedSamples   []model.Sample
+		expectedSamples   []chunks.Sample
 	}{
 		{
 			name: "In order ct+normal sample",
 			appendableSamples: []appendableSamples{
 				{ts: 100, val: 10, ct: 1},
 			},
-			expectedSamples: []model.Sample{
-				{Timestamp: 1, Value: 0},
-				{Timestamp: 100, Value: 10},
+			expectedSamples: []chunks.Sample{
+				sample{t: 1, f: 0},
+				sample{t: 100, f: 10},
 			},
 		},
 		{
@@ -5963,10 +5962,10 @@ func TestHeadAppender_AppendCTZeroSample(t *testing.T) {
 				{ts: 100, val: 10, ct: 1},
 				{ts: 101, val: 10, ct: 1},
 			},
-			expectedSamples: []model.Sample{
-				{Timestamp: 1, Value: 0},
-				{Timestamp: 100, Value: 10},
-				{Timestamp: 101, Value: 10},
+			expectedSamples: []chunks.Sample{
+				sample{t: 1, f: 0},
+				sample{t: 100, f: 10},
+				sample{t: 101, f: 10},
 			},
 		},
 		{
@@ -5975,11 +5974,11 @@ func TestHeadAppender_AppendCTZeroSample(t *testing.T) {
 				{ts: 100, val: 10, ct: 1},
 				{ts: 102, val: 10, ct: 101},
 			},
-			expectedSamples: []model.Sample{
-				{Timestamp: 1, Value: 0},
-				{Timestamp: 100, Value: 10},
-				{Timestamp: 101, Value: 0},
-				{Timestamp: 102, Value: 10},
+			expectedSamples: []chunks.Sample{
+				sample{t: 1, f: 0},
+				sample{t: 100, f: 10},
+				sample{t: 101, f: 0},
+				sample{t: 102, f: 10},
 			},
 		},
 		{
@@ -5988,41 +5987,33 @@ func TestHeadAppender_AppendCTZeroSample(t *testing.T) {
 				{ts: 100, val: 10, ct: 1},
 				{ts: 101, val: 10, ct: 100},
 			},
-			expectedSamples: []model.Sample{
-				{Timestamp: 1, Value: 0},
-				{Timestamp: 100, Value: 10},
-				{Timestamp: 101, Value: 10},
+			expectedSamples: []chunks.Sample{
+				sample{t: 1, f: 0},
+				sample{t: 100, f: 10},
+				sample{t: 101, f: 10},
 			},
 		},
 	} {
-		h, _ := newTestHead(t, DefaultBlockDuration, wlog.CompressionNone, false)
-		defer func() {
-			require.NoError(t, h.Close())
-		}()
-		a := h.Appender(context.Background())
-		lbls := labels.FromStrings("foo", "bar")
-		for _, sample := range tc.appendableSamples {
-			_, err := a.AppendCTZeroSample(0, lbls, sample.ts, sample.ct)
-			require.NoError(t, err)
-			_, err = a.Append(0, lbls, sample.ts, sample.val)
-			require.NoError(t, err)
-		}
-		require.NoError(t, a.Commit())
+		t.Run(tc.name, func(t *testing.T) {
+			h, _ := newTestHead(t, DefaultBlockDuration, wlog.CompressionNone, false)
+			defer func() {
+				require.NoError(t, h.Close())
+			}()
+			a := h.Appender(context.Background())
+			lbls := labels.FromStrings("foo", "bar")
+			for _, sample := range tc.appendableSamples {
+				_, err := a.AppendCTZeroSample(0, lbls, sample.ts, sample.ct)
+				require.NoError(t, err)
+				_, err = a.Append(0, lbls, sample.ts, sample.val)
+				require.NoError(t, err)
+			}
+			require.NoError(t, a.Commit())
 
-		q, err := NewBlockQuerier(h, math.MinInt64, math.MaxInt64)
-		require.NoError(t, err)
-		ss := q.Select(context.Background(), false, nil, labels.MustNewMatcher(labels.MatchEqual, "foo", "bar"))
-		require.True(t, ss.Next())
-		s := ss.At()
-		require.False(t, ss.Next())
-		it := s.Iterator(nil)
-		for _, sample := range tc.expectedSamples {
-			require.Equal(t, chunkenc.ValFloat, it.Next())
-			timestamp, value := it.At()
-			require.Equal(t, sample.Timestamp, model.Time(timestamp))
-			require.Equal(t, sample.Value, model.SampleValue(value))
-		}
-		require.Equal(t, chunkenc.ValNone, it.Next())
+			q, err := NewBlockQuerier(h, math.MinInt64, math.MaxInt64)
+			require.NoError(t, err)
+			result := query(t, q, labels.MustNewMatcher(labels.MatchEqual, "foo", "bar"))
+			require.Equal(t, tc.expectedSamples, result[`{foo="bar"}`])
+		})
 	}
 }