Moved labels map into helper function

Also updated tests and followed the new query name
for preemptible, resNodeLabels -> resIsSpot.

pkg/costmodel/cluster.go

@@ -500,7 +500,8 @@ func ClusterNodes(cp cloud.Provider, client prometheus.Client, duration, offset
 
 	cpuBreakdownMap := buildCPUBreakdownMap(resNodeCPUModeTotal)
 	activeDataMap := buildActiveDataMap(resActiveMins, resolution, cp.ParseID)
-	preemptibleMap := buildPreemptibleMap(resNodeLabels, cp.ParseID)
+	preemptibleMap := buildPreemptibleMap(resIsSpot, cp.ParseID)
+	labelsMap := buildLabelsMap(resLabels)
 
 	nodeMap := buildNodeMap(
 		cpuCostMap, ramCostMap, gpuCostMap,
@@ -509,31 +510,10 @@ func ClusterNodes(cp cloud.Provider, client prometheus.Client, duration, offset
 		cpuBreakdownMap,
 		activeDataMap,
 		preemptibleMap,
+		labelsMap,
 		clusterAndNameToType,
 	)
 
-	// Copy labels into node
-	for _, result := range resLabels {
-		cluster, err := result.GetString("cluster_id")
-		if err != nil {
-			cluster = env.GetClusterID()
-		}
-		node, err := result.GetString("kubernetes_node")
-		if err != nil {
-			log.DedupedWarningf(5, "ClusterNodes: label data missing node")
-			continue
-		}
-		key := fmt.Sprintf("%s/%s", cluster, node)
-		if _, ok := nodeMap[key]; !ok {
-			continue
-		}
-		for name, value := range result.Metric {
-			if val, ok := value.(string); ok {
-				nodeMap[key].Labels[name] = val
-			}
-		}
-	}
-
 	c, err := cp.GetConfig()
 	if err != nil {
 		return nil, err

pkg/costmodel/cluster_helpers.go

@@ -421,13 +421,13 @@ func buildActiveDataMap(resActiveMins []*prom.QueryResult, resolution time.Durat
 // Determine preemptibility with node labels
 // node id -> is preemptible?
 func buildPreemptibleMap(
-	resNodeLabels []*prom.QueryResult,
+	resIsSpot []*prom.QueryResult,
 	providerIDParser func(string) string,
 ) map[NodeIdentifier]bool {
 
 	m := make(map[NodeIdentifier]bool)
 
-	for _, result := range resNodeLabels {
+	for _, result := range resIsSpot {
 		nodeName, err := result.GetString("node")
 		if err != nil {
 			continue
@@ -463,6 +463,39 @@ func buildPreemptibleMap(
 	return m
 }
 
+func buildLabelsMap(
+	resLabels []*prom.QueryResult,
+) map[nodeIdentifierNoProviderID]map[string]string {
+
+	m := make(map[nodeIdentifierNoProviderID]map[string]string)
+
+	// Copy labels into node
+	for _, result := range resLabels {
+		cluster, err := result.GetString("cluster_id")
+		if err != nil {
+			cluster = env.GetClusterID()
+		}
+		node, err := result.GetString("kubernetes_node")
+		if err != nil {
+			log.DedupedWarningf(5, "ClusterNodes: label data missing node")
+			continue
+		}
+		key := nodeIdentifierNoProviderID{
+			Cluster: cluster,
+			Name:    node,
+		}
+
+		m[key] = make(map[string]string)
+
+		for name, value := range result.Metric {
+			if val, ok := value.(string); ok {
+				m[key][name] = val
+			}
+		}
+	}
+	return m
+}
+
 // checkForKeyAndInitIfMissing inits a key in the provided nodemap if
 // it does not exist. Intended to be called ONLY by buildNodeMap
 func checkForKeyAndInitIfMissing(
@@ -557,6 +590,7 @@ func buildNodeMap(
 	cpuBreakdownMap map[nodeIdentifierNoProviderID]*ClusterCostsBreakdown,
 	activeDataMap map[NodeIdentifier]activeData,
 	preemptibleMap map[NodeIdentifier]bool,
+	labelsMap map[nodeIdentifierNoProviderID]map[string]string,
 	clusterAndNameToType map[nodeIdentifierNoProviderID]string,
 ) map[NodeIdentifier]*Node {
 
@@ -619,6 +653,10 @@ func buildNodeMap(
 		if cpuBreakdown, ok := cpuBreakdownMap[clusterAndNameID]; ok {
 			nodePtr.CPUBreakdown = cpuBreakdown
 		}
+
+		if labels, ok := labelsMap[clusterAndNameID]; ok {
+			nodePtr.Labels = labels
+		}
 	}
 
 	return nodeMap

pkg/costmodel/cluster_helpers_test.go

@@ -139,6 +139,7 @@ func TestBuildNodeMap(t *testing.T) {
 		cpuBreakdownMap      map[nodeIdentifierNoProviderID]*ClusterCostsBreakdown
 		activeDataMap        map[NodeIdentifier]activeData
 		preemptibleMap       map[NodeIdentifier]bool
+		labelsMap            map[nodeIdentifierNoProviderID]map[string]string
 		clusterAndNameToType map[nodeIdentifierNoProviderID]string
 		expected             map[NodeIdentifier]*Node
 	}{
@@ -409,6 +410,22 @@ func TestBuildNodeMap(t *testing.T) {
 					ProviderID: "prov_node2_A",
 				}: false,
 			},
+			labelsMap: map[nodeIdentifierNoProviderID]map[string]string{
+				nodeIdentifierNoProviderID{
+					Cluster: "cluster1",
+					Name:    "node1",
+				}: map[string]string{
+					"labelname1_A": "labelvalue1_A",
+					"labelname1_B": "labelvalue1_B",
+				},
+				nodeIdentifierNoProviderID{
+					Cluster: "cluster1",
+					Name:    "node2",
+				}: map[string]string{
+					"labelname2_A": "labelvalue2_A",
+					"labelname2_B": "labelvalue2_B",
+				},
+			},
 			clusterAndNameToType: map[nodeIdentifierNoProviderID]string{
 				nodeIdentifierNoProviderID{
 					Cluster: "cluster1",
@@ -446,6 +463,10 @@ func TestBuildNodeMap(t *testing.T) {
 					End:         time.Date(2020, 6, 16, 9, 20, 39, 0, time.UTC),
 					Minutes:     5*60 + 35 + (11.0 / 60.0),
 					Preemptible: true,
+					Labels: map[string]string{
+						"labelname1_A": "labelvalue1_A",
+						"labelname1_B": "labelvalue1_B",
+					},
 				},
 				NodeIdentifier{
 					Cluster:    "cluster1",
@@ -473,6 +494,10 @@ func TestBuildNodeMap(t *testing.T) {
 					End:         time.Date(2020, 6, 16, 9, 21, 39, 0, time.UTC),
 					Minutes:     5*60 + 36 + (11.0 / 60.0),
 					Preemptible: false,
+					Labels: map[string]string{
+						"labelname1_A": "labelvalue1_A",
+						"labelname1_B": "labelvalue1_B",
+					},
 				},
 				NodeIdentifier{
 					Cluster:    "cluster1",
@@ -500,6 +525,10 @@ func TestBuildNodeMap(t *testing.T) {
 					End:         time.Date(2020, 6, 16, 9, 10, 39, 0, time.UTC),
 					Minutes:     5*60 + 25 + (11.0 / 60.0),
 					Preemptible: false,
+					Labels: map[string]string{
+						"labelname2_A": "labelvalue2_A",
+						"labelname2_B": "labelvalue2_B",
+					},
 				},
 			},
 		},
@@ -514,6 +543,7 @@ func TestBuildNodeMap(t *testing.T) {
 			testCase.cpuBreakdownMap,
 			testCase.activeDataMap,
 			testCase.preemptibleMap,
+			testCase.labelsMap,
 			testCase.clusterAndNameToType,
 		)