Basic functional options support to metrics assertion in integration tests

integration/alertmanager_test.go

@@ -5,8 +5,8 @@ package main
 import (
 	"context"
 	"testing"
-	"time"
 
+	"github.com/prometheus/prometheus/pkg/labels"
 	"github.com/stretchr/testify/require"
 
 	"github.com/cortexproject/cortex/integration/e2e"
@@ -80,8 +80,9 @@ func TestAlertmanagerStoreAPI(t *testing.T) {
 	err = c.SetAlertmanagerConfig(context.Background(), cortexAlertmanagerUserConfigYaml, map[string]string{})
 	require.NoError(t, err)
 
-	time.Sleep(2 * time.Second)
-	require.NoError(t, am.WaitSumMetrics(e2e.Equals(0), "cortex_alertmanager_config_invalid"))
+	require.NoError(t, am.WaitSumMetricsWithOptions(e2e.Equals(0), []string{"cortex_alertmanager_config_invalid"},
+		e2e.WithLabelMatchers(labels.MustNewMatcher(labels.MatchEqual, "user", "user-1")),
+		e2e.WaitMissingMetrics))
 
 	cfg, err := c.GetAlertmanagerConfig(context.Background())
 	require.NoError(t, err)
@@ -97,7 +98,10 @@ func TestAlertmanagerStoreAPI(t *testing.T) {
 	err = c.DeleteAlertmanagerConfig(context.Background())
 	require.NoError(t, err)
 
-	time.Sleep(2 * time.Second)
+	// The deleted config is applied asynchronously, so we should wait until the metric
+	// disappear for the specific user.
+	require.NoError(t, am.WaitMissingMetric("cortex_alertmanager_config_invalid", e2e.WithLabelMatchers(
+		labels.MustNewMatcher(labels.MatchEqual, "user", "user-1"))))
 
 	cfg, err = c.GetAlertmanagerConfig(context.Background())
 	require.Error(t, err)

integration/api_ruler_test.go

@@ -6,6 +6,7 @@ import (
 	"path/filepath"
 	"testing"
 
+	"github.com/prometheus/prometheus/pkg/labels"
 	"github.com/prometheus/prometheus/pkg/rulefmt"
 	"github.com/stretchr/testify/require"
 	"gopkg.in/yaml.v3"
@@ -137,6 +138,9 @@ func TestRulerAPISingleBinary(t *testing.T) {
 	require.Equal(t, retrievedNamespace[0].Name, "rule")
 
 	// Check to make sure prometheus engine metrics are available for both engine types
-	require.NoError(t, cortex.WaitForMetricWithLabels(e2e.EqualsSingle(0), "prometheus_engine_queries", map[string]string{"engine": "querier"}))
-	require.NoError(t, cortex.WaitForMetricWithLabels(e2e.EqualsSingle(0), "prometheus_engine_queries", map[string]string{"engine": "ruler"}))
+	require.NoError(t, cortex.WaitSumMetricsWithOptions(e2e.Equals(0), []string{"prometheus_engine_queries"}, e2e.WithLabelMatchers(
+		labels.MustNewMatcher(labels.MatchEqual, "engine", "querier"))))
+
+	require.NoError(t, cortex.WaitSumMetricsWithOptions(e2e.Equals(0), []string{"prometheus_engine_queries"}, e2e.WithLabelMatchers(
+		labels.MustNewMatcher(labels.MatchEqual, "engine", "ruler"))))
 }

integration/e2e/composite_service.go

@@ -5,6 +5,8 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/pkg/errors"
+
 	"github.com/cortexproject/cortex/pkg/util"
 )
 
@@ -39,52 +41,41 @@ func (s *CompositeHTTPService) Instances() []*HTTPService {
 // WaitSumMetrics waits for at least one instance of each given metric names to be present and their sums, returning true
 // when passed to given isExpected(...).
 func (s *CompositeHTTPService) WaitSumMetrics(isExpected func(sums ...float64) bool, metricNames ...string) error {
+	return s.WaitSumMetricsWithOptions(isExpected, metricNames)
+}
+
+func (s *CompositeHTTPService) WaitSumMetricsWithOptions(isExpected func(sums ...float64) bool, metricNames []string, opts ...MetricsOption) error {
 	var (
-		sums []float64
-		err  error
+		sums    []float64
+		err     error
+		options = buildMetricsOptions(opts)
 	)
 
 	for s.retryBackoff.Reset(); s.retryBackoff.Ongoing(); {
-		sums, err = s.SumMetrics(metricNames...)
-		if err != nil {
-			return err
-		}
-
-		if isExpected(sums...) {
-			return nil
+		sums, err = s.SumMetrics(metricNames, opts...)
+		if options.WaitMissingMetrics && errors.Is(err, errMissingMetric) {
+			continue
 		}
-
-		s.retryBackoff.Wait()
-	}
-
-	return fmt.Errorf("unable to find metrics %s with expected values. Last values: %v", metricNames, sums)
-}
-
-func (s *CompositeHTTPService) WaitSumMetricWithLabels(isExpected func(sums float64) bool, metricName string, expectedLabels map[string]string) error {
-	lastSum := 0.0
-
-	for s.retryBackoff.Reset(); s.retryBackoff.Ongoing(); {
-		lastSum, err := s.SumMetricWithLabels(metricName, expectedLabels)
 		if err != nil {
 			return err
 		}
 
-		if isExpected(lastSum) {
+		if isExpected(sums...) {
 			return nil
 		}
 
 		s.retryBackoff.Wait()
 	}
 
-	return fmt.Errorf("unable to find metric %s with labels %v with expected value. Last value: %v", metricName, expectedLabels, lastSum)
+	return fmt.Errorf("unable to find metrics %s with expected values. Last error: %v. Last values: %v", metricNames, err, sums)
 }
 
 // SumMetrics returns the sum of the values of each given metric names.
-func (s *CompositeHTTPService) SumMetrics(metricNames ...string) ([]float64, error) {
+func (s *CompositeHTTPService) SumMetrics(metricNames []string, opts ...MetricsOption) ([]float64, error) {
 	sums := make([]float64, len(metricNames))
 
 	for _, service := range s.services {
-		partials, err := service.SumMetrics(metricNames...)
+		partials, err := service.SumMetrics(metricNames, opts...)
 		if err != nil {
 			return nil, err
 		}
@@ -100,19 +91,3 @@ func (s *CompositeHTTPService) SumMetrics(metricNames ...string) ([]float64, err
 
 	return sums, nil
 }
-
-// SumMetricWithLabels returns the sum of the values of metric with matching labels across all services.
-func (s *CompositeHTTPService) SumMetricWithLabels(metricName string, expectedLabels map[string]string) (float64, error) {
-	sum := 0.0
-
-	for _, service := range s.services {
-		s, err := service.SumMetricWithLabels(metricName, expectedLabels)
-		if err != nil {
-			return 0, err
-		}
-
-		sum += s
-	}
-
-	return sum, nil
-}

integration/e2e/metrics.go

@@ -6,7 +6,7 @@ import (
 	io_prometheus_client "github.com/prometheus/client_model/go"
 )
 
-func getValue(m *io_prometheus_client.Metric) float64 {
+func getMetricValue(m *io_prometheus_client.Metric) float64 {
 	if m.GetGauge() != nil {
 		return m.GetGauge().GetValue()
 	} else if m.GetCounter() != nil {
@@ -20,10 +20,63 @@ func getValue(m *io_prometheus_client.Metric) float64 {
 	}
 }
 
-func sumValues(family *io_prometheus_client.MetricFamily) float64 {
+func getMetricCount(m *io_prometheus_client.Metric) float64 {
+	if m.GetHistogram() != nil {
+		return float64(m.GetHistogram().GetSampleCount())
+	} else if m.GetSummary() != nil {
+		return float64(m.GetSummary().GetSampleCount())
+	} else {
+		return 0
+	}
+}
+
+func getValues(metrics []*io_prometheus_client.Metric, opts MetricsOptions) []float64 {
+	values := make([]float64, 0, len(metrics))
+	for _, m := range metrics {
+		values = append(values, opts.GetValue(m))
+	}
+	return values
+}
+
+func filterMetrics(metrics []*io_prometheus_client.Metric, opts MetricsOptions) []*io_prometheus_client.Metric {
+	// If no label matcher is configured, then no filtering should be done.
+	if len(opts.LabelMatchers) == 0 {
+		return metrics
+	}
+	if len(metrics) == 0 {
+		return metrics
+	}
+
+	filtered := make([]*io_prometheus_client.Metric, 0, len(metrics))
+
+	for _, m := range metrics {
+		metricLabels := map[string]string{}
+		for _, lp := range m.GetLabel() {
+			metricLabels[lp.GetName()] = lp.GetValue()
+		}
+
+		matches := true
+		for _, matcher := range opts.LabelMatchers {
+			if !matcher.Matches(metricLabels[matcher.Name]) {
+				matches = false
+				break
+			}
+		}
+
+		if !matches {
+			continue
+		}
+
+		filtered = append(filtered, m)
+	}
+
+	return filtered
+}
+
+func sumValues(values []float64) float64 {
 	sum := 0.0
-	for _, m := range family.Metric {
-		sum += getValue(m)
+	for _, v := range values {
+		sum += v
 	}
 	return sum
 }

integration/e2e/metrics_options.go

@@ -0,0 +1,52 @@
+package e2e
+
+import (
+	io_prometheus_client "github.com/prometheus/client_model/go"
+	"github.com/prometheus/prometheus/pkg/labels"
+)
+
+var (
+	DefaultMetricsOptions = MetricsOptions{
+		GetValue:           getMetricValue,
+		WaitMissingMetrics: false,
+	}
+)
+
+// GetMetricValueFunc defined the signature of a function used to get the metric value.
+type GetMetricValueFunc func(m *io_prometheus_client.Metric) float64
+
+// MetricsOption defined the signature of a function used to manipulate options.
+type MetricsOption func(*MetricsOptions)
+
+// MetricsOptions is the structure holding all options.
+type MetricsOptions struct {
+	GetValue           GetMetricValueFunc
+	LabelMatchers      []*labels.Matcher
+	WaitMissingMetrics bool
+}
+
+// WithMetricCount is an option to get the histogram/summary count as metric value.
+func WithMetricCount(opts *MetricsOptions) {
+	opts.GetValue = getMetricCount
+}
+
+// WithLabelMatchers is an option to filter only matching series.
+func WithLabelMatchers(matchers ...*labels.Matcher) MetricsOption {
+	return func(opts *MetricsOptions) {
+		opts.LabelMatchers = matchers
+	}
+}
+
+// WithWaitMissingMetrics is an option to wait whenever an expected metric is missing. If this
+// option is not enabled, will return error on missing metrics.
+func WaitMissingMetrics(opts *MetricsOptions) {
+	opts.WaitMissingMetrics = true
+}
+
+func buildMetricsOptions(opts []MetricsOption) MetricsOptions {
+	result := DefaultMetricsOptions
+	for _, opt := range opts {
+		opt(&result)
+	}
+	return result
+}