Fix rField fallback. Add more tests.

Signed-off-by: Utku Ozdemir <uoz@protonmail.com>
echoblag · Jun 24, 2021 · 2963abd · 2963abd
1 parent 0324c0f
commit 2963abd
Show file tree

Hide file tree

Showing 2 changed files with 99 additions and 7 deletions.
diff --git a/internal/exporter/exporter.go b/internal/exporter/exporter.go
@@ -72,7 +72,7 @@ func buildQFieldToRFieldMap(logger log.Logger, qFieldsRaw string,
 	nvidiaSmiCommand string) (map[qField]rField, error) {
 	qFieldsSeparated := strings.Split(qFieldsRaw, ",")
 
-	var parsedQFields []qField
+	qFields := toQFieldSlice(qFieldsSeparated)
 	if len(qFieldsSeparated) == 1 && qFieldsSeparated[0] == qFieldsAuto {
 		parsed, err := ParseAutoQFields(nvidiaSmiCommand)
 		if err != nil {
@@ -82,17 +82,25 @@ func buildQFieldToRFieldMap(logger log.Logger, qFieldsRaw string,
 			return fallbackQFieldToRFieldMap, nil
 		}
 
-		parsedQFields = parsed
+		qFields = parsed
 	}
 
-	t, err := scrape(parsedQFields, nvidiaSmiCommand)
+	t, err := scrape(qFields, nvidiaSmiCommand)
+	var rFields []rField
 	if err != nil {
-		return nil, err
+		_ = level.Warn(logger).Log("msg",
+			"Failed to run an initial scrape, using the built-in list for field mapping")
+		rFields, err = getFallbackValues(qFields)
+		if err != nil {
+			return nil, err
+		}
+	} else {
+		rFields = t.rFields
 	}
 
-	r := make(map[qField]rField, len(parsedQFields))
-	for i, qField := range parsedQFields {
-		r[qField] = t.rFields[i]
+	r := make(map[qField]rField, len(qFields))
+	for i, qField := range qFields {
+		r[qField] = rFields[i]
 	}
 
 	return r, nil
@@ -253,3 +261,18 @@ func getKeys(m map[qField]rField) []qField {
 	}
 	return r
 }
+
+func getFallbackValues(qFields []qField) ([]rField, error) {
+	r := make([]rField, len(qFields))
+	i := 0
+	for _, qField := range qFields {
+		val, contains := fallbackQFieldToRFieldMap[qField]
+		if !contains {
+			return nil, fmt.Errorf("unexpected query field: %s", qField)
+		}
+
+		r[i] = val
+		i++
+	}
+	return r, nil
+}
diff --git a/internal/exporter/exporter_test.go b/internal/exporter/exporter_test.go
@@ -1,8 +1,10 @@
 package exporter
 
 import (
+	"github.com/go-kit/log"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/stretchr/testify/assert"
+	"strings"
 	"testing"
 )
 
@@ -116,3 +118,70 @@ func TestBuildQFieldToMetricInfoMap(t *testing.T) {
 	assertFloat(t, 1, metricInfo2.valueMultiplier)
 	assert.Equal(t, prometheus.GaugeValue, metricInfo2.mType)
 }
+
+func TestNewUnknownField(t *testing.T) {
+	logger := log.NewNopLogger()
+	_, err := New("aaa", "bbb", "a", logger)
+	assert.Error(t, err)
+}
+
+func TestDescribe(t *testing.T) {
+	logger := log.NewNopLogger()
+	exp, err := New("aaa", "bbb", "fan.speed,memory.used", logger)
+	assert.NoError(t, err)
+
+	doneCh := make(chan bool)
+	descCh := make(chan *prometheus.Desc)
+
+	go func() {
+		exp.Describe(descCh)
+		doneCh <- true
+	}()
+
+	var descStrs []string
+
+end:
+	for {
+		select {
+		case desc := <-descCh:
+			descStrs = append(descStrs, desc.String())
+		case <-doneCh:
+			break end
+		}
+	}
+
+	assert.Len(t, descStrs, 3)
+	descs := strings.Join(descStrs, "\n")
+	assert.Contains(t, descs, "aaa_fan_speed")
+	assert.Contains(t, descs, "aaa_memory_used")
+	assert.Contains(t, descs, "aaa_failed_scrapes_total")
+}
+
+func TestCollect(t *testing.T) {
+	logger := log.NewNopLogger()
+	exp, err := New("aaa", "bbb", "fan.speed,memory.used", logger)
+	assert.NoError(t, err)
+
+	doneCh := make(chan bool)
+	metricCh := make(chan prometheus.Metric)
+
+	go func() {
+		exp.Collect(metricCh)
+		doneCh <- true
+	}()
+
+	var metrics []string
+
+end:
+	for {
+		select {
+		case metric := <-metricCh:
+			metrics = append(metrics, metric.Desc().String())
+		case <-doneCh:
+			break end
+		}
+	}
+
+	assert.Len(t, metrics, 1)
+	assert.Contains(t, metrics[0], "aaa_failed_scrapes_total")
+}