fix(teststorage/appender.go): TODO and Sample staleness check (#17905)

* fix(teststorage/appender.go): TODO and Sample staleness check

Allow different order of consecutive stale samples between the expected
and actual array for RequireEqual and RequireNotEqual by trying to
swap the expected side until it matches.

Also fix the definition of stale sample in the test, it's not only
float, but defined for native histograms as well.

Signed-off-by: György Krajcsovits <gyorgy.krajcsovits@grafana.com>

* add unit tests

Signed-off-by: György Krajcsovits <gyorgy.krajcsovits@grafana.com>

---------

Signed-off-by: György Krajcsovits <gyorgy.krajcsovits@grafana.com>

util/teststorage/appender.go

@@ -24,7 +24,6 @@ import (
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
-	"github.com/google/go-cmp/cmp/cmpopts"
 	"github.com/prometheus/common/model"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/atomic"
@@ -97,37 +96,32 @@ func (s Sample) Equals(other Sample) bool {
 		slices.EqualFunc(s.ES, other.ES, exemplar.Exemplar.Equals)
 }
 
-var (
-	sampleComparer = cmp.Comparer(func(a, b Sample) bool {
-		return a.Equals(b)
-	})
-	byLabelSort = cmpopts.SortSlices(func(a, b Sample) int {
-		return labels.Compare(a.L, b.L)
-	})
-)
-
-func includeStaleNaNs(s []Sample) bool {
-	for _, e := range s {
-		if value.IsStaleNaN(e.V) {
-			return true
-		}
+// IsStale returns whether the sample represents a stale sample, according to
+// https://prometheus.io/docs/specs/native_histograms/#staleness-markers.
+func (s Sample) IsStale() bool {
+	switch {
+	case s.FH != nil:
+		return value.IsStaleNaN(s.FH.Sum)
+	case s.H != nil:
+		return value.IsStaleNaN(s.H.Sum)
+	default:
+		return value.IsStaleNaN(s.V)
 	}
-	return false
 }
 
+var sampleComparer = cmp.Comparer(func(a, b Sample) bool {
+	return a.Equals(b)
+})
+
 // RequireEqual is a special require equal that correctly compare Prometheus structures.
 //
 // In comparison to testutil.RequireEqual, this function adds special logic for comparing []Samples.
 //
-// It also ignores ordering when expected slice contains at least one StaleNaN. This is because the
-// scrape StaleNan samples are generated by iterating over a map, thus expectedly different.
-//
-// TODO(bwplotka): We should likely reorder only within a group of sequential NaNs or only in scrape package.
+// It also ignores ordering between consecutive stale samples to avoid false
+// negatives due to map iteration order in staleness tracking.
 func RequireEqual(t testing.TB, expected, got []Sample, msgAndArgs ...any) {
 	opts := []cmp.Option{sampleComparer}
-	if includeStaleNaNs(expected) {
-		opts = append(opts, byLabelSort)
-	}
+	expected = reorderExpectedForStaleness(expected, got)
 	testutil.RequireEqualWithOptions(t, expected, got, opts, msgAndArgs...)
 }
 
@@ -136,9 +130,7 @@ func RequireNotEqual(t testing.TB, expected, got []Sample, msgAndArgs ...any) {
 	t.Helper()
 
 	opts := []cmp.Option{cmp.Comparer(labels.Equal), sampleComparer}
-	if includeStaleNaNs(expected) {
-		opts = append(opts, byLabelSort)
-	}
+	expected = reorderExpectedForStaleness(expected, got)
 	if !cmp.Equal(expected, got, opts...) {
 		return
 	}
@@ -147,6 +139,45 @@ func RequireNotEqual(t testing.TB, expected, got []Sample, msgAndArgs ...any) {
 		"b: %s", expected, got), msgAndArgs...)
 }
 
+func reorderExpectedForStaleness(expected, got []Sample) []Sample {
+	if len(expected) != len(got) || !includeStaleNaNs(expected) {
+		return expected
+	}
+	result := make([]Sample, len(expected))
+	copy(result, expected)
+
+	// Try to reorder only consecutive stale samples to avoid false negatives
+	// due to map iteration order in staleness tracking.
+	for i := range result {
+		if !result[i].IsStale() {
+			continue
+		}
+		if result[i].Equals(got[i]) {
+			continue
+		}
+		for j := i + 1; j < len(result); j++ {
+			if !result[j].IsStale() {
+				break
+			}
+			if result[j].Equals(got[i]) {
+				// Swap.
+				result[i], result[j] = result[j], result[i]
+				break
+			}
+		}
+	}
+	return result
+}
+
+func includeStaleNaNs(s []Sample) bool {
+	for _, e := range s {
+		if e.IsStale() {
+			return true
+		}
+	}
+	return false
+}
+
 // Appendable is a storage.Appendable mock.
 // It allows recording all samples that were added through the appender and injecting errors.
 // Appendable will panic if more than one Appender is open.

util/teststorage/appender_test.go

@@ -306,3 +306,108 @@ func TestConcurrentAppenderV2_ReturnsErrAppender(t *testing.T) {
 	require.Error(t, app.Commit())
 	require.Error(t, app.Rollback())
 }
+
+func TestReorderExpectedForStaleness(t *testing.T) {
+	testcases := []struct {
+		name       string
+		inExpected []Sample
+		inGot      []Sample
+		expected   []Sample
+	}{
+		{
+			name: "no staleness markers",
+			inExpected: []Sample{
+				{L: labels.FromStrings("a", "1"), T: 1, V: 1},
+				{L: labels.FromStrings("a", "2"), T: 1, V: 2},
+			},
+			inGot: []Sample{
+				{L: labels.FromStrings("a", "2"), T: 1, V: 2},
+				{L: labels.FromStrings("a", "1"), T: 1, V: 1},
+			},
+		},
+		{
+			name: "with staleness markers",
+			inExpected: []Sample{
+				{L: labels.FromStrings("a", "1"), T: 1, V: 1},
+				{L: labels.FromStrings("a", "2"), T: 2, V: 2},
+				{L: labels.FromStrings("a", "3"), T: 3, V: math.Float64frombits(value.StaleNaN)},
+				{L: labels.FromStrings("a", "4"), T: 4, V: math.Float64frombits(value.StaleNaN)},
+			},
+			inGot: []Sample{
+				{L: labels.FromStrings("a", "1"), T: 1, V: 1},
+				{L: labels.FromStrings("a", "2"), T: 2, V: 2},
+				{L: labels.FromStrings("a", "3"), T: 3, V: math.Float64frombits(value.StaleNaN)},
+				{L: labels.FromStrings("a", "4"), T: 4, V: math.Float64frombits(value.StaleNaN)},
+			},
+		},
+		{
+			name: "with staleness markers wrong order",
+			inExpected: []Sample{
+				{L: labels.FromStrings("a", "1"), T: 1, V: 1},
+				{L: labels.FromStrings("a", "2"), T: 2, V: 2},
+				{L: labels.FromStrings("a", "3"), T: 3, V: math.Float64frombits(value.StaleNaN)},
+				{L: labels.FromStrings("a", "4"), T: 4, V: math.Float64frombits(value.StaleNaN)},
+			},
+			inGot: []Sample{
+				{L: labels.FromStrings("a", "2"), T: 2, V: 2},
+				{L: labels.FromStrings("a", "1"), T: 1, V: 1},
+				{L: labels.FromStrings("a", "4"), T: 4, V: math.Float64frombits(value.StaleNaN)},
+				{L: labels.FromStrings("a", "3"), T: 3, V: math.Float64frombits(value.StaleNaN)},
+			},
+			expected: []Sample{
+				{L: labels.FromStrings("a", "1"), T: 1, V: 1},
+				{L: labels.FromStrings("a", "2"), T: 2, V: 2},
+				{L: labels.FromStrings("a", "4"), T: 4, V: math.Float64frombits(value.StaleNaN)},
+				{L: labels.FromStrings("a", "3"), T: 3, V: math.Float64frombits(value.StaleNaN)},
+			},
+		},
+		{
+			name: "with staleness markers wrong order but not consecutive",
+			inExpected: []Sample{
+				{L: labels.FromStrings("a", "1"), T: 1, V: 1},
+				{L: labels.FromStrings("a", "3"), T: 3, V: math.Float64frombits(value.StaleNaN)},
+				{L: labels.FromStrings("a", "2"), T: 2, V: 2},
+				{L: labels.FromStrings("a", "4"), T: 4, V: math.Float64frombits(value.StaleNaN)},
+			},
+			inGot: []Sample{
+				{L: labels.FromStrings("a", "2"), T: 2, V: 2},
+				{L: labels.FromStrings("a", "1"), T: 1, V: 1},
+				{L: labels.FromStrings("a", "4"), T: 4, V: math.Float64frombits(value.StaleNaN)},
+				{L: labels.FromStrings("a", "3"), T: 3, V: math.Float64frombits(value.StaleNaN)},
+			},
+			expected: []Sample{
+				{L: labels.FromStrings("a", "1"), T: 1, V: 1},
+				{L: labels.FromStrings("a", "3"), T: 3, V: math.Float64frombits(value.StaleNaN)},
+				{L: labels.FromStrings("a", "2"), T: 2, V: 2},
+				{L: labels.FromStrings("a", "4"), T: 4, V: math.Float64frombits(value.StaleNaN)},
+			},
+		},
+	}
+	for _, tc := range testcases {
+		t.Run(tc.name, func(t *testing.T) {
+			if tc.expected == nil {
+				tc.expected = tc.inExpected
+			}
+			RequireEqual(t, tc.expected, reorderExpectedForStaleness(tc.inExpected, tc.inGot))
+		})
+	}
+}
+
+func TestSampleIsStale(t *testing.T) {
+	s1 := Sample{V: 1}
+	require.False(t, s1.IsStale())
+	s2 := Sample{V: math.Float64frombits(value.StaleNaN)}
+	require.True(t, s2.IsStale())
+	h := tsdbutil.GenerateTestHistogram(0)
+	h1 := Sample{V: math.Float64frombits(value.StaleNaN), H: h}
+	require.False(t, h1.IsStale()) // Histogram takes precedence over V.
+	h.Sum = math.Float64frombits(value.StaleNaN)
+	h2 := Sample{V: 1, H: h}
+	require.True(t, h2.IsStale())
+	fh := tsdbutil.GenerateTestFloatHistogram(0)
+	fh1 := Sample{V: math.Float64frombits(value.StaleNaN), H: h, FH: fh}
+	require.False(t, fh1.IsStale()) // FloatHistogram takes precedence over all.
+	fh.Sum = math.Float64frombits(value.StaleNaN)
+	fh2 := Sample{V: 1, H: tsdbutil.GenerateTestHistogram(1), FH: fh}
+	require.True(t, fh2.IsStale())
+}