Alerting: Fix private labels filtering test (#109393)

pkg/services/ngalert/schedule/recording_rule_test.go

@@ -12,8 +12,11 @@ import (
 	"time"
 
 	"github.com/benbjohnson/clock"
+	"github.com/gogo/protobuf/proto"
+	"github.com/golang/snappy"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/testutil"
+	"github.com/prometheus/prometheus/prompb"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/atomic"
 
@@ -756,12 +759,12 @@ func testRecordingRule_Integration(t *testing.T, writeTarget *writer.TestRemoteW
 		rule := gen.With(withQueryForHealth("ok")).GenerateRef()
 		rule.Record.TargetDatasourceUID = dsUID
 		rule.Labels = map[string]string{
-			"normal_label":                             "value1",
+			"normal_label":                             "not_filtered_1",
-			"another_label":                            "value2",
+			"another_label":                            "not_filtered_2",
 			models.AutogeneratedRouteLabel:             "filtered",
 			models.AutogeneratedRouteReceiverNameLabel: "filtered",
-			"__user_custom__":                          "not_filtered",
+			"__user_custom__":                          "not_filtered_3",
-			"only_end__":                               "not_filtered",
+			"only_end__":                               "not_filtered_4",
 		}
 
 		ruleStore.PutRule(context.Background(), rule)
@@ -790,12 +793,27 @@ func testRecordingRule_Integration(t *testing.T, writeTarget *writer.TestRemoteW
 			require.Equal(t, 1, writeTarget.RequestsCount)
 			require.NotEmpty(t, writeTarget.LastRequestBody)
 
-			// Check that the body doesn't contain the private labels
+			writeReq := decodePrometheusWriteRequest(t, writeTarget.LastRequestBody)
-			require.NotContains(t, writeTarget.LastRequestBody, models.AutogeneratedRouteLabel)
-			require.NotContains(t, writeTarget.LastRequestBody, models.AutogeneratedRouteReceiverNameLabel)
 
-			require.Contains(t, writeTarget.LastRequestBody, "__user_custom__")
+			// Check that the private labels are filtered out
-			require.Contains(t, writeTarget.LastRequestBody, rule.Record.Metric)
+			for _, label := range []string{
+				models.AutogeneratedRouteLabel,
+				models.AutogeneratedRouteReceiverNameLabel,
+			} {
+				_, found := getLabel(writeReq, label)
+				require.False(t, found, "Label %s should not be present in the write request", label)
+			}
+
+			// Check that user-defined labels are preserved
+			for _, label := range []string{
+				"normal_label",
+				"another_label",
+				"__user_custom__",
+				"only_end__",
+			} {
+				value, _ := getLabel(writeReq, label)
+				require.Equal(t, rule.Labels[label], value, "Label %s=%s should be present in the write request", label, rule.Labels[label])
+			}
 		})
 	})
 }
@@ -874,3 +892,28 @@ type mockPluginContextProvider struct{}
 func (m *mockPluginContextProvider) GetWithDataSource(ctx context.Context, pluginID string, user identity.Requester, ds *datasources.DataSource) (backend.PluginContext, error) {
 	return backend.PluginContext{}, nil
 }
+
+func decodePrometheusWriteRequest(t *testing.T, data string) *prompb.WriteRequest {
+	t.Helper()
+
+	decompressed, err := snappy.Decode(nil, []byte(data))
+	require.NoError(t, err)
+
+	var writeReq prompb.WriteRequest
+	err = proto.Unmarshal(decompressed, &writeReq)
+	require.NoError(t, err)
+
+	return &writeReq
+}
+
+func getLabel(req *prompb.WriteRequest, labelName string) (string, bool) {
+	for _, ts := range req.Timeseries {
+		for _, label := range ts.Labels {
+			if label.Name == labelName {
+				return label.Value, true
+			}
+		}
+	}
+
+	return "", false
+}