From 0577ef6c57494933d4139f19dd2960359736ad1e Mon Sep 17 00:00:00 2001
From: beorn7 <beorn@grafana.com>
Date: Thu, 14 May 2020 00:17:45 +0200
Subject: [PATCH] Improve CollectAndCount

Now that we have also added CollectAndLint and GatherAndLint, I
thought we should bring CollectAndCount in line. So:

- Add GatherAndCount.
- Add filtering by metric name.
- Add tests.

Minor wart: CollectAndCount should now return `(int, error)`, but that
would be a breaking change as the current version just returns
`int`. I decided to let the new version panic when it should return an
error. An error is anyway very unlikely, so the biggest annoyance here
is really just the inconsistency.

Signed-off-by: beorn7 <beorn@grafana.com>
---
 prometheus/testutil/testutil.go      | 58 +++++++++++++++++-----------
 prometheus/testutil/testutil_test.go | 27 +++++++++++++
 2 files changed, 62 insertions(+), 23 deletions(-)

diff --git a/prometheus/testutil/testutil.go b/prometheus/testutil/testutil.go
index c47373c..9af60ce 100644
--- a/prometheus/testutil/testutil.go
+++ b/prometheus/testutil/testutil.go
@@ -112,31 +112,43 @@ func ToFloat64(c prometheus.Collector) float64 {
 	panic(fmt.Errorf("collected a non-gauge/counter/untyped metric: %s", pb))
 }
 
-// CollectAndCount collects all Metrics from the provided Collector and returns their number.
-//
-// This can be used to assert the number of metrics collected by a given collector after certain operations.
-//
-// This function is only for testing purposes, and even for testing, other approaches
-// are often more appropriate (see this package's documentation).
-func CollectAndCount(c prometheus.Collector) int {
-	var (
-		mCount int
-		mChan  = make(chan prometheus.Metric)
-		done   = make(chan struct{})
-	)
+// CollectAndCount registers the provided Collector with a newly created
+// pedantic Registry. It then calls GatherAndCount with that Registry and with
+// the provided metricNames. In the unlikely case that the registration or the
+// gathering fails, this function panics. (This is inconsistent with the other
+// CollectAnd… functions in this package and has historical reasons. Changing
+// the function signature would be a breaking change and will therefore only
+// happen with the next major version bump.)
+func CollectAndCount(c prometheus.Collector, metricNames ...string) int {
+	reg := prometheus.NewPedanticRegistry()
+	if err := reg.Register(c); err != nil {
+		panic(fmt.Errorf("registering collector failed: %s", err))
+	}
+	result, err := GatherAndCount(reg, metricNames...)
+	if err != nil {
+		panic(err)
+	}
+	return result
+}
 
-	go func() {
-		for range mChan {
-			mCount++
-		}
-		close(done)
-	}()
+// GatherAndCount gathers all metrics from the provided Gatherer and counts
+// them. It returns the number of metric children in all gathered metric
+// families together. If any metricNames are provided, only metrics with those
+// names are counted.
+func GatherAndCount(g prometheus.Gatherer, metricNames ...string) (int, error) {
+	got, err := g.Gather()
+	if err != nil {
+		return 0, fmt.Errorf("gathering metrics failed: %s", err)
+	}
+	if metricNames != nil {
+		got = filterMetrics(got, metricNames)
+	}
 
-	c.Collect(mChan)
-	close(mChan)
-	<-done
-
-	return mCount
+	result := 0
+	for _, mf := range got {
+		result += len(mf.GetMetric())
+	}
+	return result, nil
 }
 
 // CollectAndCompare registers the provided Collector with a newly created
diff --git a/prometheus/testutil/testutil_test.go b/prometheus/testutil/testutil_test.go
index aaf6707..56d9933 100644
--- a/prometheus/testutil/testutil_test.go
+++ b/prometheus/testutil/testutil_test.go
@@ -306,3 +306,30 @@ some_total{label1="value1"} 1
 		t.Errorf("Expected\n%#+v\nGot:\n%#+v", expectedError, err.Error())
 	}
 }
+
+func TestCollectAndCount(t *testing.T) {
+	c := prometheus.NewCounterVec(
+		prometheus.CounterOpts{
+			Name: "some_total",
+			Help: "A value that represents a counter.",
+		},
+		[]string{"foo"},
+	)
+	if got, want := CollectAndCount(c), 0; got != want {
+		t.Errorf("unexpected metric count, got %d, want %d", got, want)
+	}
+	c.WithLabelValues("bar")
+	if got, want := CollectAndCount(c), 1; got != want {
+		t.Errorf("unexpected metric count, got %d, want %d", got, want)
+	}
+	c.WithLabelValues("baz")
+	if got, want := CollectAndCount(c), 2; got != want {
+		t.Errorf("unexpected metric count, got %d, want %d", got, want)
+	}
+	if got, want := CollectAndCount(c, "some_total"), 2; got != want {
+		t.Errorf("unexpected metric count, got %d, want %d", got, want)
+	}
+	if got, want := CollectAndCount(c, "some_other_total"), 0; got != want {
+		t.Errorf("unexpected metric count, got %d, want %d", got, want)
+	}
+}