Simplify ExtractValueFromMetric function

xonvanetta · Mar 24, 2023 · 0a5dc83 · 0a5dc83
1 parent a32ca15
commit 0a5dc83
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 57 deletions.
diff --git a/prometheus_helper/prometheus_helper.go b/prometheus_helper/prometheus_helper.go
@@ -42,56 +42,21 @@ func ExtractFqName(metric string) (string, error) {
 	return matches[1], nil
 }
 
-// CustomCollector is a custom prometheus.Collector that collects only the given metric.
-type CustomCollector struct {
-	metric prometheus.Metric
-}
-
-// Describe implements the prometheus.Collector interface.
-func (c *CustomCollector) Describe(ch chan<- *prometheus.Desc) {
-	ch <- c.metric.Desc()
-}
-
-// Collect implements the prometheus.Collector interface.
-func (c *CustomCollector) Collect(ch chan<- prometheus.Metric) {
-	ch <- c.metric
-}
-
 // ExtractValueFromMetric extracts the value from a prometheus.Metric object.
 // It creates a custom collector and registry, registers the given metric, and then collects
 // the metric value using the registry.
 // Returns the extracted float64 value from the metric's Gauge.
 func ExtractValueFromMetric(metric prometheus.Metric) (float64, error) {
-	// Custom collector that collects only the given metric.
-	collector := &CustomCollector{
-		metric: metric,
-	}
-
-	// Create a custom registry and register the collector.
-	registry := prometheus.NewRegistry()
-	err := registry.Register(collector)
+	var dtoMetric dto.Metric
+	err := metric.Write(&dtoMetric)
 	if err != nil {
-		return 0, err
+		return 0, fmt.Errorf("error writing metric: %v", err)
 	}
 
-	var metricValue float64
-	metricChannel := make(chan prometheus.Metric)
-	go func() {
-		registry.Collect(metricChannel)
-		close(metricChannel)
-	}()
-
-	for collectedMetric := range metricChannel {
-		if collectedMetric.Desc().String() == metric.Desc().String() {
-			var dtoMetric dto.Metric
-			err = collectedMetric.Write(&dtoMetric)
-			if err != nil {
-				return 0, fmt.Errorf("error writing metric: %v", err)
-			}
-			metricValue = dtoMetric.GetGauge().GetValue()
-			break
-		}
+	gauge := dtoMetric.GetGauge()
+	if gauge == nil {
+		return 0, errors.New("the metric is not a Gauge")
 	}
 
-	return metricValue, nil
+	return gauge.GetValue(), nil
 }
diff --git a/prometheus_helper/prometheus_helper_test.go b/prometheus_helper/prometheus_helper_test.go
@@ -1,10 +1,12 @@
 package prometheus_helper
 
 import (
+	"errors"
 	"fmt"
 	"testing"
 
 	"github.com/prometheus/client_golang/prometheus"
+	dto "github.com/prometheus/client_model/go"
 )
 
 func TestSimpleDescHelper(t *testing.T) {
@@ -65,6 +67,16 @@ func TestExtractFqdnName(t *testing.T) {
 	})
 }
 
+type badMetricStub struct{}
+
+func (m *badMetricStub) Desc() *prometheus.Desc {
+	return nil
+}
+
+func (m *badMetricStub) Write(*dto.Metric) error {
+	return errors.New("writing metric failed")
+}
+
 func TestExtractValueFromMetric(t *testing.T) {
 	t.Run("should extract value from a metric", func(t *testing.T) {
 		metricDesc := prometheus.NewDesc("test_metric", "test metric help", nil, nil)
@@ -81,30 +93,31 @@ func TestExtractValueFromMetric(t *testing.T) {
 		}
 	})
 
-	t.Run("should return an error when unable to write metric", func(t *testing.T) {
-		metricDesc := prometheus.NewDesc("test_metric", "test metric help", nil, nil)
-		exampleErr := fmt.Errorf("example error")
-		invalidMetric := prometheus.NewInvalidMetric(metricDesc, exampleErr)
+	t.Run("should return error if writing metric fails", func(t *testing.T) {
+		badMetric := &badMetricStub{}
+		val, err := ExtractValueFromMetric(badMetric)
 
-		customCollector := &CustomCollector{
-			metric: invalidMetric,
+		if err == nil {
+			t.Errorf("Expected error, but got nil")
 		}
 
-		registry := prometheus.NewRegistry()
-		err := registry.Register(customCollector)
-
-		if err != nil {
-			t.Errorf("Unexpected error: %v", err)
+		if val != 0 {
+			t.Errorf("Expected value to be 0, got %f", val)
 		}
+	})
 
-		extractedValue, err := ExtractValueFromMetric(invalidMetric)
+	t.Run("should return error if the metric is not a Gauge", func(t *testing.T) {
+		metricDesc := prometheus.NewDesc("test_counter_metric", "test counter metric help", nil, nil)
+		metricValue := 42.0
+		metric := prometheus.MustNewConstMetric(metricDesc, prometheus.CounterValue, metricValue)
 
+		val, err := ExtractValueFromMetric(metric)
 		if err == nil {
-			t.Errorf("Expected error but got nil")
+			t.Errorf("Expected error, but got nil")
 		}
 
-		if extractedValue != 0 {
-			t.Errorf("Expected extracted value to be 0, got %f", extractedValue)
+		if val != 0 {
+			t.Errorf("Expected value to be 0, got %f", val)
 		}
 	})
 }