cloudprober · manugarg · May 10, 2023 · May 10, 2023 · May 10, 2023 · May 10, 2023
diff --git a/metrics/testutils/testutils.go b/metrics/testutils/testutils.go
@@ -47,26 +47,76 @@ func MetricsFromChannel(dataChan chan *metrics.EventMetrics, num int, timeout ti
 	return
 }
 
-// MetricsMap rearranges a list of metrics into a map of map.
-// {
-//   "m1": {
-//       "target1": [val1, val2..],
-//       "target2": [val1],
-//   },
-//   "m2": {
-//       ...
-//   }
-// }
-func MetricsMap(ems []*metrics.EventMetrics) map[string]map[string][]*metrics.EventMetrics {
-	results := make(map[string]map[string][]*metrics.EventMetrics)
+func LabelsMapByTarget(ems []*metrics.EventMetrics) map[string]map[string]string {
+	lmap := make(map[string]map[string]string)
 	for _, em := range ems {
 		target := em.Label("dst")
-		for _, m := range em.MetricsKeys() {
-			if results[m] == nil {
-				results[m] = make(map[string][]*metrics.EventMetrics)
+		if lmap[target] == nil {
+			lmap[target] = make(map[string]string)
+		}
+		for _, k := range em.LabelsKeys() {
+			lmap[target][k] = em.Label(k)
+		}
+	}
+	return lmap
+}
+
+type MetricsMap map[string]map[string][]metrics.Value
+
+// MetricsMapByTarget rearranges a list of metrics into a map of map.
+//
+//	{
+//	  "target1": {
+//	      "m1": [val1, val2..],
+//	      "m2": [val1],
+//	  },
+//	  "target2": {
+//	      ...
+//	  }
+//	}
+func MetricsMapByTarget(ems []*metrics.EventMetrics) MetricsMap {
+	mmap := make(map[string]map[string][]metrics.Value)
+	for _, em := range ems {
+		target := em.Label("dst")
+		if mmap[target] == nil {
+			mmap[target] = make(map[string][]metrics.Value)
+		}
+		for _, k := range em.MetricsKeys() {
+			mmap[target][k] = append(mmap[target][k], em.Metric(k))
+		}
+	}
+	return mmap
+}
+
+// Filter returns a map of target to values for a given metric name.
+func (mmap MetricsMap) Filter(metricName string) map[string][]metrics.Value {
+	res := make(map[string][]metrics.Value, len(mmap))
+	for tgt, valMap := range mmap {
+		for name, vals := range valMap {
+			if name != metricName {
+				continue
+			}
+			res[tgt] = vals
+		}
+	}
+	return res
+}
+
+// LastValueInt64 returns the last value for a given metric name and target.
+func (mmap MetricsMap) LastValueInt64(target, metricName string) int64 {
+	for tgt, valMap := range mmap {
+		if tgt != target {
+			continue
+		}
+		for name, vals := range valMap {
+			if name != metricName {
+				continue
+			}
+			if len(vals) == 0 {
+				return -1
 			}
-			results[m][target] = append(results[m][target], em)
+			return vals[len(vals)-1].(metrics.NumValue).Int64()
 		}
 	}
-	return results
+	return -1
 }
diff --git a/metrics/testutils/testutils_test.go b/metrics/testutils/testutils_test.go
@@ -19,6 +19,7 @@ import (
 	"time"
 
 	"github.com/cloudprober/cloudprober/metrics"
+	"github.com/stretchr/testify/assert"
 )
 
 func TestMetricsFromChannel(t *testing.T) {
@@ -64,8 +65,8 @@ func TestMetricsFromChannel(t *testing.T) {
 func TestMetricsMap(t *testing.T) {
 	var ems []*metrics.EventMetrics
 	expectedValues := map[string][]int64{
-		"success": []int64{99, 98},
-		"total":   []int64{100, 100},
+		"success": {99, 98},
+		"total":   {100, 100},
 	}
 	ems = append(ems, metrics.NewEventMetrics(time.Now()).
 		AddMetric("success", metrics.NewInt(99)).
@@ -76,23 +77,68 @@ func TestMetricsMap(t *testing.T) {
 		AddMetric("total", metrics.NewInt(100)).
 		AddLabel("dst", "target2"))
 
-	metricsMap := MetricsMap(ems)
+	metricsMap := MetricsMapByTarget(ems)
 
-	for _, m := range []string{"success", "total"} {
-		if metricsMap[m] == nil {
-			t.Errorf("didn't get metric %s in metrics map", m)
+	for i, tgt := range []string{"target1", "target2"} {
+		for _, m := range []string{"success", "total"} {
+			assert.Len(t, metricsMap[tgt][m], 1, "number of values mismatch")
+			val := metricsMap[tgt][m][0].(metrics.NumValue).Int64()
+			assert.Equal(t, expectedValues[m][i], val, "metric value mismatch")
 		}
 	}
+}
 
-	for i, tgt := range []string{"target1", "target2"} {
-		for _, m := range []string{"success", "total"} {
-			if len(metricsMap[m][tgt]) != 1 {
-				t.Errorf("Wrong number of values for metric (%s) for target (%s) from the command output. Got=%d, Expected=1", m, tgt, len(metricsMap[m][tgt]))
-			}
-			val := metricsMap[m][tgt][0].Metric(m).(metrics.NumValue).Int64()
-			if val != expectedValues[m][i] {
-				t.Errorf("Wrong metric value for target (%s) from the command output. Got=%d, Expected=%d", m, val, expectedValues[m][i])
+func TestMetricsMapFilterAndLastValueInt64(t *testing.T) {
+	mmap := MetricsMap{
+		"target1": map[string][]metrics.Value{
+			"success": {metrics.NewInt(98), metrics.NewInt(99)},
+			"total":   {metrics.NewInt(100), metrics.NewInt(101)},
+		},
+		"target2": map[string][]metrics.Value{
+			"success2": {metrics.NewInt(198), metrics.NewInt(199)},
+			"total":    {metrics.NewInt(200), metrics.NewInt(201)},
+		},
+	}
+
+	tests := []struct {
+		name       string
+		target     string
+		metricName string
+		filtered   map[string][]metrics.Value
+		want       int64
+	}{
+		{
+			name:       "target2",
+			target:     "target1",
+			metricName: "success",
+			filtered: map[string][]metrics.Value{
+				"target1": {metrics.NewInt(98), metrics.NewInt(99)},
+			},
+			want: 99,
+		},
+		{
+			name:       "target2",
+			target:     "target2",
+			metricName: "success",
+			filtered: map[string][]metrics.Value{
+				"target1": {metrics.NewInt(98), metrics.NewInt(99)},
+			},
+			want: -1,
+		},
+		{
+			name:       "not found",
+			target:     "target1",
+			metricName: "latency",
+			filtered:   map[string][]metrics.Value{},
+			want:       -1,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert.Equal(t, tt.filtered, mmap.Filter(tt.metricName), "filtered map mismatch")
+			if got := mmap.LastValueInt64(tt.target, tt.metricName); got != tt.want {
+				t.Errorf("MetricsMap.LastValueInt64() = %v, want %v", got, tt.want)
 			}
-		}
+		})
 	}
 }
diff --git a/probes/common/sched/sched_test.go b/probes/common/sched/sched_test.go
@@ -39,9 +39,9 @@ func (tpr *testProbeResult) Metrics(ts time.Time, opts *options.Options) *metric
 func compareNumberOfMetrics(t *testing.T, ems []*metrics.EventMetrics, metricName string, targets [2]string, wantCloseRange bool) {
 	t.Helper()
 
-	m := testutils.MetricsMap(ems)[metricName]
-	num1 := len(m[targets[0]])
-	num2 := len(m[targets[1]])
+	mmap := testutils.MetricsMapByTarget(ems).Filter(metricName)
+	num1 := len(mmap[targets[0]])
+	num2 := len(mmap[targets[1]])
 
 	diff := num1 - num2
 	threshold := num1 / 2
@@ -58,9 +58,9 @@ func compareNumberOfMetrics(t *testing.T, ems []*metrics.EventMetrics, metricNam
 	// than (if stats_export_interval > interval), or equal to the number of
 	// metrics received. Note: This test assumes that metric value is
 	// incremented in each run.
-	for _, ems := range m {
-		numMetrics := len(ems)
-		lastVal := int(ems[numMetrics-1].Metric(metricName).(metrics.NumValue).Int64())
+	for _, mvs := range mmap {
+		numMetrics := len(mvs)
+		lastVal := int(mvs[numMetrics-1].(metrics.NumValue).Int64())
 		if lastVal < numMetrics {
 			t.Errorf("Metric (%s) last value: %d, less than: %d", metricName, lastVal, numMetrics)
 		}