From 36b9f46811b73b12b09ca9b6cf65bc04291db555 Mon Sep 17 00:00:00 2001
From: Jack Cassidy <j.cassidy45@gmail.com>
Date: Thu, 9 May 2024 16:08:21 +0100
Subject: [PATCH] add: CollectAndFormat to testutil (#1503)

CollectAndFormat is a helper function that returns the formatted metrics
to the caller, allowing them to use it how they want. This is different
to CollectAndCompare where the comparison is done strictly on behalf of
the caller. Often it is more convenient to perform a simple substring
match on the formatted metric.

Signed-off-by: Jack Cassidy <j.cassidy45@gmail.com>
---
 prometheus/testutil/testutil.go      | 30 ++++++++++++++++++++++++---
 prometheus/testutil/testutil_test.go | 31 ++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 3 deletions(-)

diff --git a/prometheus/testutil/testutil.go b/prometheus/testutil/testutil.go
index 7c1e905..e818c2f 100644
--- a/prometheus/testutil/testutil.go
+++ b/prometheus/testutil/testutil.go
@@ -183,9 +183,8 @@ func ScrapeAndCompare(url string, expected io.Reader, metricNames ...string) err
 	return compareMetricFamilies(scraped, wanted, metricNames...)
 }
 
-// CollectAndCompare registers the provided Collector with a newly created
-// pedantic Registry. It then calls GatherAndCompare with that Registry and with
-// the provided metricNames.
+// CollectAndCompare collects the metrics identified by `metricNames` and compares them in the Prometheus text
+// exposition format to the data read from expected.
 func CollectAndCompare(c prometheus.Collector, expected io.Reader, metricNames ...string) error {
 	reg := prometheus.NewPedanticRegistry()
 	if err := reg.Register(c); err != nil {
@@ -221,6 +220,31 @@ func TransactionalGatherAndCompare(g prometheus.TransactionalGatherer, expected
 	return compareMetricFamilies(got, wanted, metricNames...)
 }
 
+// CollectAndFormat collects the metrics identified by `metricNames` and returns them in the given format.
+func CollectAndFormat(c prometheus.Collector, format expfmt.FormatType, metricNames ...string) ([]byte, error) {
+	reg := prometheus.NewPedanticRegistry()
+	if err := reg.Register(c); err != nil {
+		return nil, fmt.Errorf("registering collector failed: %w", err)
+	}
+
+	gotFiltered, err := reg.Gather()
+	if err != nil {
+		return nil, fmt.Errorf("gathering metrics failed: %w", err)
+	}
+
+	gotFiltered = filterMetrics(gotFiltered, metricNames)
+
+	var gotFormatted bytes.Buffer
+	enc := expfmt.NewEncoder(&gotFormatted, expfmt.NewFormat(format))
+	for _, mf := range gotFiltered {
+		if err := enc.Encode(mf); err != nil {
+			return nil, fmt.Errorf("encoding gathered metrics failed: %w", err)
+		}
+	}
+
+	return gotFormatted.Bytes(), nil
+}
+
 // convertReaderToMetricFamily would read from a io.Reader object and convert it to a slice of
 // dto.MetricFamily.
 func convertReaderToMetricFamily(reader io.Reader) ([]*dto.MetricFamily, error) {
diff --git a/prometheus/testutil/testutil_test.go b/prometheus/testutil/testutil_test.go
index 40374a0..06e3677 100644
--- a/prometheus/testutil/testutil_test.go
+++ b/prometheus/testutil/testutil_test.go
@@ -20,6 +20,8 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/prometheus/common/expfmt"
+
 	"github.com/prometheus/client_golang/prometheus"
 )
 
@@ -431,3 +433,32 @@ func TestCollectAndCount(t *testing.T) {
 		t.Errorf("unexpected metric count, got %d, want %d", got, want)
 	}
 }
+
+func TestCollectAndFormat(t *testing.T) {
+	const expected = `# HELP foo_bar A value that represents the number of bars in foo.
+# TYPE foo_bar counter
+foo_bar{fizz="bang"} 1
+`
+	c := prometheus.NewCounterVec(
+		prometheus.CounterOpts{
+			Name: "foo_bar",
+			Help: "A value that represents the number of bars in foo.",
+		},
+		[]string{"fizz"},
+	)
+	c.WithLabelValues("bang").Inc()
+
+	got, err := CollectAndFormat(c, expfmt.TypeTextPlain, "foo_bar")
+	if err != nil {
+		t.Errorf("unexpected error: %s", err.Error())
+	}
+
+	gotS := string(got)
+	if err != nil {
+		t.Errorf("unexpected error: %s", err.Error())
+	}
+
+	if gotS != expected {
+		t.Errorf("unexpected metric output, got %q, expected %q", gotS, expected)
+	}
+}