From e185de152c89c723a31b93856de014ae4beb24b8 Mon Sep 17 00:00:00 2001
From: Bartlomiej Plotka <bwplotka@gmail.com>
Date: Thu, 9 Jan 2020 10:07:07 +0000
Subject: [PATCH] Added testutil.MetricCount

Signed-off-by: Bartlomiej Plotka <bwplotka@gmail.com>
---
 prometheus/testutil/testutil.go | 28 ++++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/prometheus/testutil/testutil.go b/prometheus/testutil/testutil.go
index 4c72bc6cb..1108b451d 100644
--- a/prometheus/testutil/testutil.go
+++ b/prometheus/testutil/testutil.go
@@ -108,6 +108,34 @@ func ToFloat64(c prometheus.Collector) float64 {
 	panic(fmt.Errorf("collected a non-gauge/counter/untyped metric: %s", pb))
 }
 
+// MetricCount collects all Metrics from the provided Collector and counts them.
+//
+// Typical use case is when you want to assert certain number of series of given collector after
+// some operation.
+//
+// This function is only for testing purposes, and even for testing, other approaches
+// are often more appropriate (see this package's documentation).
+func MetricCount(c prometheus.Collector) int {
+	var (
+		mCount int
+		mChan  = make(chan prometheus.Metric)
+		done   = make(chan struct{})
+	)
+
+	go func() {
+		for range mChan {
+			mCount++
+		}
+		close(done)
+	}()
+
+	c.Collect(mChan)
+	close(mChan)
+	<-done
+
+	return mCount
+}
+
 // CollectAndCompare registers the provided Collector with a newly created
 // pedantic Registry. It then does the same as GatherAndCompare, gathering the
 // metrics from the pedantic Registry.