From 5ba0993f6f91a7be093f0460572ee88d435cc94c Mon Sep 17 00:00:00 2001
From: beorn7 <beorn@soundcloud.com>
Date: Thu, 23 Aug 2018 00:21:38 +0200
Subject: [PATCH] Improved interface

- Expected text format is now read from an io.Reader.
- Metrics are gathered from a Gatherer.
- Added a convenience wrapper to collect from a Collector.

Signed-off-by: beorn7 <beorn@soundcloud.com>
---
 prometheus/testutil/testutil.go      | 29 ++++++++++++++++------------
 prometheus/testutil/testutil_test.go | 11 ++++++-----
 2 files changed, 23 insertions(+), 17 deletions(-)

diff --git a/prometheus/testutil/testutil.go b/prometheus/testutil/testutil.go
index 6145833..5d731ae 100644
--- a/prometheus/testutil/testutil.go
+++ b/prometheus/testutil/testutil.go
@@ -16,6 +16,7 @@ package testutil
 import (
 	"bytes"
 	"fmt"
+	"io"
 	"reflect"
 	"sort"
 
@@ -26,15 +27,23 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 )
 
-// GatherAndCompare retrieves all metrics exposed by a collector and compares it
-// to an expected output in the Prometheus text exposition format.
-// metricNames allows only comparing the given metrics. All are compared if it's nil.
-func GatherAndCompare(c prometheus.Collector, expected string, metricNames ...string) error {
+// 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.
+func CollectAndCompare(c prometheus.Collector, expected io.Reader, metricNames ...string) error {
 	reg := prometheus.NewPedanticRegistry()
 	if err := reg.Register(c); err != nil {
 		return fmt.Errorf("registering collector failed: %s", err)
 	}
-	metrics, err := reg.Gather()
+	return GatherAndCompare(reg, expected, metricNames...)
+}
+
+// GatherAndCompare gathers all metrics from the provided Gatherer and compares
+// it to an expected output read from the provided Reader in the Prometheus text
+// exposition format. If any metricNames are provided, only metrics with those
+// names are compared.
+func GatherAndCompare(g prometheus.Gatherer, expected io.Reader, metricNames ...string) error {
+	metrics, err := g.Gather()
 	if err != nil {
 		return fmt.Errorf("gathering metrics failed: %s", err)
 	}
@@ -42,13 +51,13 @@ func GatherAndCompare(c prometheus.Collector, expected string, metricNames ...st
 		metrics = filterMetrics(metrics, metricNames)
 	}
 	var tp expfmt.TextParser
-	expectedMetrics, err := tp.TextToMetricFamilies(bytes.NewReader([]byte(expected)))
+	expectedMetrics, err := tp.TextToMetricFamilies(expected)
 	if err != nil {
 		return fmt.Errorf("parsing expected metrics failed: %s", err)
 	}
 
 	if !reflect.DeepEqual(metrics, normalizeMetricFamilies(expectedMetrics)) {
-		// Encode the gathered output to the readbale text format for comparison.
+		// Encode the gathered output to the readable text format for comparison.
 		var buf1 bytes.Buffer
 		enc := expfmt.NewEncoder(&buf1, expfmt.FmtText)
 		for _, mf := range metrics {
@@ -82,16 +91,12 @@ got:
 func filterMetrics(metrics []*dto.MetricFamily, names []string) []*dto.MetricFamily {
 	var filtered []*dto.MetricFamily
 	for _, m := range metrics {
-		drop := true
 		for _, name := range names {
 			if m.GetName() == name {
-				drop = false
+				filtered = append(filtered, m)
 				break
 			}
 		}
-		if !drop {
-			filtered = append(filtered, m)
-		}
 	}
 	return filtered
 }
diff --git a/prometheus/testutil/testutil_test.go b/prometheus/testutil/testutil_test.go
index a05eda7..b4e7b01 100644
--- a/prometheus/testutil/testutil_test.go
+++ b/prometheus/testutil/testutil_test.go
@@ -1,4 +1,4 @@
-// Copyright 2018 he Prometheus Authors
+// Copyright 2018 The Prometheus Authors
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
@@ -14,12 +14,13 @@
 package testutil
 
 import (
+	"strings"
 	"testing"
 
 	"github.com/prometheus/client_golang/prometheus"
 )
 
-func TestGatherAndCompare(t *testing.T) {
+func TestCollectAndCompare(t *testing.T) {
 	const metadata = `
 		# HELP some_total A value that represents a counter.
 		# TYPE some_total counter
@@ -39,7 +40,7 @@ func TestGatherAndCompare(t *testing.T) {
 		some_total{ label1 = "value1" } 1
 	`
 
-	if err := GatherAndCompare(c, metadata+expected, "some_total"); err != nil {
+	if err := CollectAndCompare(c, strings.NewReader(metadata+expected), "some_total"); err != nil {
 		t.Errorf("unexpected collecting result:\n%s", err)
 	}
 }
@@ -63,7 +64,7 @@ func TestNoMetricFilter(t *testing.T) {
 		some_total{label1="value1"} 1
 	`
 
-	if err := GatherAndCompare(c, metadata+expected); err != nil {
+	if err := CollectAndCompare(c, strings.NewReader(metadata+expected)); err != nil {
 		t.Errorf("unexpected collecting result:\n%s", err)
 	}
 }
@@ -103,7 +104,7 @@ some_total{label1="value1"} 1
 
 `
 
-	err := GatherAndCompare(c, metadata+expected)
+	err := CollectAndCompare(c, strings.NewReader(metadata+expected))
 	if err == nil {
 		t.Error("Expected error, got no error.")
 	}