Add WriteToTextfile test

Signed-off-by: Sevag Hanssian <sevag.hanssian@gmail.com>

prometheus/registry.go

@@ -16,6 +16,7 @@ package prometheus
 import (
 	"bytes"
 	"fmt"
+	"os"
 	"runtime"
 	"sort"
 	"strconv"
@@ -539,10 +540,12 @@ func (r *Registry) WriteToTextfile(path string) error {
 	if err != nil {
 		return err
 	}
-	output := []string{}
+	outputMap := map[string][]string{}
+	metricNames := []string{}
 	for _, metricFamily := range metricFamilies {
+		output := []string{}
 		output = append(output, fmt.Sprintf("# HELP %s %s", metricFamily.GetName(), metricFamily.GetHelp()))
-		output = append(output, fmt.Sprintf("# TYPE %s %s", metricFamily.GetName(), metricFamily.GetType().String()))
+		output = append(output, fmt.Sprintf("# TYPE %s %s", metricFamily.GetName(), strings.ToLower(metricFamily.GetType().String())))
 		for _, metric := range metricFamily.GetMetric() {
 			labelStrings := []string{}
 			if metric.GetLabel() != nil {
@@ -570,33 +573,61 @@ func (r *Registry) WriteToTextfile(path string) error {
 			case dto.MetricType_SUMMARY:
 				labelString := fmt.Sprintf("{%s}", strings.Join(labelStrings, ","))
 				count := metric.GetSummary().GetSampleCount()
-				output = append(output, fmt.Sprintf("%s_count%s %d%s", metricFamily.GetName(), labelString, count, timestampString))
 				sum := strconv.FormatFloat(metric.GetSummary().GetSampleSum(), 'f', -1, 64)
-				output = append(output, fmt.Sprintf("%s_sum%s %s%s", metricFamily.GetName(), labelString, sum, timestampString))
 				for _, quantile := range metric.GetSummary().GetQuantile() {
 					quantileName := strconv.FormatFloat(quantile.GetQuantile(), 'f', -1, 64)
 					quantileLabelStrings := append(labelStrings, fmt.Sprintf("quantile=\"%s\"", quantileName))
-					labelString = fmt.Sprintf("{%s}", strings.Join(quantileLabelStrings, ","))
+					loopLabelString := fmt.Sprintf("{%s}", strings.Join(quantileLabelStrings, ","))
 					value := strconv.FormatFloat(quantile.GetValue(), 'f', -1, 64)
-					output = append(output, fmt.Sprintf("%s_quantile%s %s%s", metricFamily.GetName(), labelString, value, timestampString))
+					output = append(output, fmt.Sprintf("%s%s %s%s", metricFamily.GetName(), loopLabelString, value, timestampString))
 				}
+				output = append(output, fmt.Sprintf("%s_sum%s %s%s", metricFamily.GetName(), labelString, sum, timestampString))
+				output = append(output, fmt.Sprintf("%s_count%s %d%s", metricFamily.GetName(), labelString, count, timestampString))
 			case dto.MetricType_HISTOGRAM:
 				labelString := fmt.Sprintf("{%s}", strings.Join(labelStrings, ","))
 				count := metric.GetHistogram().GetSampleCount()
-				output = append(output, fmt.Sprintf("%s_count%s %f%s", metricFamily.GetName(), labelString, float64(count), timestampString))
+				sum := strconv.FormatFloat(metric.GetHistogram().GetSampleSum(), 'f', -1, 64)
-				sum := metric.GetHistogram().GetSampleSum()
-				output = append(output, fmt.Sprintf("%s_sum%s %f%s", metricFamily.GetName(), labelString, sum, timestampString))
 				for _, bucket := range metric.GetHistogram().GetBucket() {
 					bucketUpperBound := strconv.FormatFloat(bucket.GetUpperBound(), 'f', -1, 64)
 					bucketLabelStrings := append(labelStrings, fmt.Sprintf("le=\"%s\"", bucketUpperBound))
-					labelString = fmt.Sprintf("{%s}", strings.Join(bucketLabelStrings, ","))
+					loopLabelString := fmt.Sprintf("{%s}", strings.Join(bucketLabelStrings, ","))
 					value := bucket.GetCumulativeCount()
-					output = append(output, fmt.Sprintf("%s_bucket%s %d%s", metricFamily.GetName(), labelString, value, timestampString))
+					output = append(output, fmt.Sprintf("%s_bucket%s %d%s", metricFamily.GetName(), loopLabelString, value, timestampString))
+				}
+				infBucketLabelStrings := append(labelStrings, "le=\"+Inf\"")
+				infLabelString := fmt.Sprintf("{%s}", strings.Join(infBucketLabelStrings, ","))
+				output = append(output, fmt.Sprintf("%s_bucket%s %d%s", metricFamily.GetName(), infLabelString, count, timestampString))
+				output = append(output, fmt.Sprintf("%s_sum%s %s%s", metricFamily.GetName(), labelString, sum, timestampString))
+				output = append(output, fmt.Sprintf("%s_count%s %d%s", metricFamily.GetName(), labelString, count, timestampString))
 			}
 		}
+		outputMap[metricFamily.GetName()] = output
+		metricNames = append(metricNames, metricFamily.GetName())
 	}
+
+	tmppath := fmt.Sprintf("%s.%d", path, os.Getpid())
+
+	f, err := os.Create(tmppath)
+	if err != nil {
+		return err
 	}
-	fmt.Println(strings.Join(output, "\n"))
+
+	sort.Strings(metricNames)
+
+	outStr := ""
+	for _, metricName := range metricNames {
+		outStr = fmt.Sprintf("%s%s\n", outStr, strings.Join(outputMap[metricName], "\n"))
+	}
+
+	_, err = f.WriteString(outStr)
+	if err != nil {
+		return err
+	}
+
+	if err := os.Rename(tmppath, path); err != nil {
+		return err
+	}
+
 	return nil
 }
 

prometheus/registry_test.go

@@ -21,9 +21,11 @@ package prometheus_test
 
 import (
 	"bytes"
+	"io/ioutil"
 	"math/rand"
 	"net/http"
 	"net/http/httptest"
+	"os"
 	"sync"
 	"testing"
 	"time"
@@ -871,3 +873,100 @@ func TestHistogramVecRegisterGatherConcurrency(t *testing.T) {
 	close(quit)
 	wg.Wait()
 }
+
+func TestWriteToTextfile(t *testing.T) {
+	expectedOut := `# HELP test_counter test counter
+# TYPE test_counter counter
+test_counter{name="qux"} 1
+# HELP test_gauge test gauge
+# TYPE test_gauge gauge
+test_gauge{name="baz"} 1.1
+# HELP test_hist test histogram
+# TYPE test_hist histogram
+test_hist_bucket{name="bar",le="0.005"} 0
+test_hist_bucket{name="bar",le="0.01"} 0
+test_hist_bucket{name="bar",le="0.025"} 0
+test_hist_bucket{name="bar",le="0.05"} 0
+test_hist_bucket{name="bar",le="0.1"} 0
+test_hist_bucket{name="bar",le="0.25"} 0
+test_hist_bucket{name="bar",le="0.5"} 0
+test_hist_bucket{name="bar",le="1"} 1
+test_hist_bucket{name="bar",le="2.5"} 1
+test_hist_bucket{name="bar",le="5"} 2
+test_hist_bucket{name="bar",le="10"} 2
+test_hist_bucket{name="bar",le="+Inf"} 2
+test_hist_sum{name="bar"} 3.64
+test_hist_count{name="bar"} 2
+# HELP test_summary test summary
+# TYPE test_summary summary
+test_summary{name="foo",quantile="0.5"} 10
+test_summary{name="foo",quantile="0.9"} 20
+test_summary{name="foo",quantile="0.99"} 20
+test_summary_sum{name="foo"} 30
+test_summary_count{name="foo"} 2
+`
+
+	registry := prometheus.NewRegistry()
+
+	summary := prometheus.NewSummaryVec(
+		prometheus.SummaryOpts{
+			Name: "test_summary",
+			Help: "test summary",
+		},
+		[]string{"name"},
+	)
+
+	histogram := prometheus.NewHistogramVec(
+		prometheus.HistogramOpts{
+			Name: "test_hist",
+			Help: "test histogram",
+		},
+		[]string{"name"},
+	)
+
+	gauge := prometheus.NewGaugeVec(
+		prometheus.GaugeOpts{
+			Name: "test_gauge",
+			Help: "test gauge",
+		},
+		[]string{"name"},
+	)
+
+	counter := prometheus.NewCounterVec(
+		prometheus.CounterOpts{
+			Name: "test_counter",
+			Help: "test counter",
+		},
+		[]string{"name"},
+	)
+
+	registry.MustRegister(summary)
+	registry.MustRegister(histogram)
+	registry.MustRegister(gauge)
+	registry.MustRegister(counter)
+
+	summary.With(prometheus.Labels{"name": "foo"}).Observe(10)
+	summary.With(prometheus.Labels{"name": "foo"}).Observe(20)
+	histogram.With(prometheus.Labels{"name": "bar"}).Observe(0.93)
+	histogram.With(prometheus.Labels{"name": "bar"}).Observe(2.71)
+	gauge.With(prometheus.Labels{"name": "baz"}).Set(1.1)
+	counter.With(prometheus.Labels{"name": "qux"}).Inc()
+
+	tmpfile, err := ioutil.TempFile("", "prom_registry_test")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.Remove(tmpfile.Name())
+
+	registry.WriteToTextfile(tmpfile.Name())
+
+	fileBytes, err := ioutil.ReadFile(tmpfile.Name())
+	if err != nil {
+		t.Fatal(err)
+	}
+	fileContents := string(fileBytes)
+
+	if fileContents != expectedOut {
+		t.Error("file contents didn't match unexpected")
+	}
+}