change metric name only if withUnit opted in

Signed-off-by: Arianna Vespri <[email protected]>

Author: vesari

expfmt/openmetrics_create.go

@@ -111,6 +111,9 @@ func MetricFamilyToOpenMetrics(out io.Writer, in *dto.MetricFamily, options ...T
 	if metricType == dto.MetricType_COUNTER && strings.HasSuffix(shortName, "_total") {
 		shortName = name[:len(name)-6]
 	}
+	if toOM.withUnit && !strings.HasSuffix(shortName, fmt.Sprintf("_%s", *in.Unit)) {
+		shortName = shortName + fmt.Sprintf("_%s", *in.Unit)
+	}
 
 	// Comments, first HELP, then TYPE.
 	if in.Help != nil {
@@ -183,6 +186,7 @@ func MetricFamilyToOpenMetrics(out io.Writer, in *dto.MetricFamily, options ...T
 		if err != nil {
 			return
 		}
+
 		err = w.WriteByte(' ')
 		written++
 		if err != nil {

expfmt/openmetrics_create_test.go

@@ -33,8 +33,9 @@ func TestCreateOpenMetrics(t *testing.T) {
 	}
 
 	var scenarios = []struct {
-		in  *dto.MetricFamily
-		out string
+		in       *dto.MetricFamily
+		withUnit ToOpenMetricsOption
+		out      string
 	}{
 		// 0: Counter, timestamp given, no _total suffix.
 		{
@@ -430,7 +431,6 @@ foos_total 42.0
 			},
 			out: `# HELP name_seconds doc string
 # TYPE name_seconds counter
-# UNIT name_seconds seconds
 `,
 		},
 		// 10: Histogram plus unit.
@@ -473,7 +473,6 @@ foos_total 42.0
 			},
 			out: `# HELP request_duration_microseconds The response latency.
 # TYPE request_duration_microseconds histogram
-# UNIT request_duration_microseconds microseconds
 request_duration_microseconds_bucket{le="100.0"} 123
 request_duration_microseconds_bucket{le="120.0"} 412
 request_duration_microseconds_bucket{le="144.0"} 592
@@ -487,6 +486,7 @@ request_duration_microseconds_count 2693
 
 	for i, scenario := range scenarios {
 		out := bytes.NewBuffer(make([]byte, 0, len(scenario.out)))
+
 		n, err := MetricFamilyToOpenMetrics(out, scenario.in)
 		if err != nil {
 			t.Errorf("%d. error: %s", i, err)
@@ -508,6 +508,69 @@ request_duration_microseconds_count 2693
 
 }
 
+func TestCreateOpenMatricsWithUnit(t *testing.T) {
+	openMetricsTimestamp := timestamppb.New(time.Unix(12345, 600000000))
+	if err := openMetricsTimestamp.CheckValid(); err != nil {
+		t.Error(err)
+	}
+
+	var scenarios = []struct {
+		in       *dto.MetricFamily
+		withUnit ToOpenMetricsOption
+		out      string
+	}{
+		// 1: No metric plus unit.
+		{
+			in: &dto.MetricFamily{
+				Name:   proto.String("name_seconds_total"),
+				Help:   proto.String("doc string"),
+				Type:   dto.MetricType_COUNTER.Enum(),
+				Unit:   proto.String("seconds"),
+				Metric: []*dto.Metric{},
+			},
+			out: `# HELP name_seconds doc string
+# TYPE name_seconds counter
+# UNIT name_seconds seconds
+`,
+		},
+		// 2: No metric plus unit no unit in name.
+		{
+			in: &dto.MetricFamily{
+				Name:   proto.String("name_total"),
+				Help:   proto.String("doc string"),
+				Type:   dto.MetricType_COUNTER.Enum(),
+				Unit:   proto.String("seconds"),
+				Metric: []*dto.Metric{},
+			},
+			out: `# HELP name_seconds doc string
+# TYPE name_seconds counter
+# UNIT name_seconds seconds
+`,
+		},
+	}
+	for i, scenario := range scenarios {
+		out := bytes.NewBuffer(make([]byte, 0, len(scenario.out)))
+
+		n, err := MetricFamilyToOpenMetrics(out, scenario.in, ToOpenMetricsWithUnit())
+		if err != nil {
+			t.Errorf("%d. error: %s", i, err)
+			continue
+		}
+		if expected, got := len(scenario.out), n; expected != got {
+			t.Errorf(
+				"%d. expected %d bytes written, got %d",
+				i, expected, got,
+			)
+		}
+		if expected, got := scenario.out, out.String(); expected != got {
+			t.Errorf(
+				"%d. expected out=%q, got %q",
+				i, expected, got,
+			)
+		}
+	}
+}
+
 func BenchmarkOpenMetricsCreate(b *testing.B) {
 	mf := &dto.MetricFamily{
 		Name: proto.String("request_duration_microseconds"),