Unit tests for E2 manual cost adjustment (#676)

* e2 fixes

* Added tests for CPU cost adjustment for e2

Co-authored-by: Ajay Tripathy <4tripathy@gmail.com>

pkg/costmodel/cluster.go

@@ -394,6 +394,9 @@ type Node struct {
 	Labels       map[string]string
 }
 
+// GKE lies about the number of cores e2 nodes have. This table
+// contains a mapping from node type -> actual CPU cores
+// for those cases.
 var partialCPUMap = map[string]float64{
 	"e2-micro":  0.25,
 	"e2-small":  0.5,

pkg/costmodel/cluster_helpers.go

@@ -187,19 +187,7 @@ func buildCPUCoresMap(
 			Cluster: cluster,
 			Name:    name,
 		}
-		if nodeType, ok := clusterAndNameToType[key]; ok {
-			if v, ok := partialCPUMap[nodeType]; ok {
-				m[key] = v
-				if cpuCores > 0 {
-					adjustmentFactor := v / cpuCores
-					m[key] = m[key] * adjustmentFactor
-				}
-			} else {
-				m[key] = cpuCores
-			}
-		} else {
-			m[key] = cpuCores
-		}
+		m[key] = cpuCores
 	}
 
 	return m
@@ -636,6 +624,13 @@ func buildNodeMap(
 
 		if cores, ok := cpuCoresMap[clusterAndNameID]; ok {
 			nodePtr.CPUCores = cores
+			if v, ok := partialCPUMap[nodePtr.NodeType]; ok {
+				if cores > 0 {
+					nodePtr.CPUCores = v
+					adjustmentFactor := v / cores
+					nodePtr.CPUCost = nodePtr.CPUCost * adjustmentFactor
+				}
+			}
 		}
 
 		if ramBytes, ok := ramBytesMap[clusterAndNameID]; ok {

pkg/costmodel/cluster_helpers_test.go

@@ -532,6 +532,120 @@ func TestBuildNodeMap(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "e2-micro cpu cost adjustment",
+			cpuCostMap: map[NodeIdentifier]float64{
+				NodeIdentifier{
+					Cluster:    "cluster1",
+					Name:       "node1",
+					ProviderID: "prov_node1",
+				}: 0.048,
+			},
+			cpuCoresMap: map[nodeIdentifierNoProviderID]float64{
+				nodeIdentifierNoProviderID{
+					Cluster: "cluster1",
+					Name:    "node1",
+				}: 6.0, // GKE lies about number of cores
+			},
+			clusterAndNameToType: map[nodeIdentifierNoProviderID]string{
+				nodeIdentifierNoProviderID{
+					Cluster: "cluster1",
+					Name:    "node1",
+				}: "e2-micro", // for this node type
+			},
+			expected: map[NodeIdentifier]*Node{
+				NodeIdentifier{
+					Cluster:    "cluster1",
+					Name:       "node1",
+					ProviderID: "prov_node1",
+				}: &Node{
+					Cluster:      "cluster1",
+					Name:         "node1",
+					ProviderID:   "prov_node1",
+					NodeType:     "e2-micro",
+					CPUCost:      0.048 * (partialCPUMap["e2-micro"] / 6.0), // adjustmentFactor is (v / GKE cores)
+					CPUCores:     partialCPUMap["e2-micro"],
+					CPUBreakdown: &ClusterCostsBreakdown{},
+					RAMBreakdown: &ClusterCostsBreakdown{},
+				},
+			},
+		},
+		{
+			name: "e2-small cpu cost adjustment",
+			cpuCostMap: map[NodeIdentifier]float64{
+				NodeIdentifier{
+					Cluster:    "cluster1",
+					Name:       "node1",
+					ProviderID: "prov_node1",
+				}: 0.048,
+			},
+			cpuCoresMap: map[nodeIdentifierNoProviderID]float64{
+				nodeIdentifierNoProviderID{
+					Cluster: "cluster1",
+					Name:    "node1",
+				}: 6.0, // GKE lies about number of cores
+			},
+			clusterAndNameToType: map[nodeIdentifierNoProviderID]string{
+				nodeIdentifierNoProviderID{
+					Cluster: "cluster1",
+					Name:    "node1",
+				}: "e2-small", // for this node type
+			},
+			expected: map[NodeIdentifier]*Node{
+				NodeIdentifier{
+					Cluster:    "cluster1",
+					Name:       "node1",
+					ProviderID: "prov_node1",
+				}: &Node{
+					Cluster:      "cluster1",
+					Name:         "node1",
+					ProviderID:   "prov_node1",
+					NodeType:     "e2-small",
+					CPUCost:      0.048 * (partialCPUMap["e2-small"] / 6.0), // adjustmentFactor is (v / GKE cores)
+					CPUCores:     partialCPUMap["e2-small"],
+					CPUBreakdown: &ClusterCostsBreakdown{},
+					RAMBreakdown: &ClusterCostsBreakdown{},
+				},
+			},
+		},
+		{
+			name: "e2-medium cpu cost adjustment",
+			cpuCostMap: map[NodeIdentifier]float64{
+				NodeIdentifier{
+					Cluster:    "cluster1",
+					Name:       "node1",
+					ProviderID: "prov_node1",
+				}: 0.048,
+			},
+			cpuCoresMap: map[nodeIdentifierNoProviderID]float64{
+				nodeIdentifierNoProviderID{
+					Cluster: "cluster1",
+					Name:    "node1",
+				}: 6.0, // GKE lies about number of cores
+			},
+			clusterAndNameToType: map[nodeIdentifierNoProviderID]string{
+				nodeIdentifierNoProviderID{
+					Cluster: "cluster1",
+					Name:    "node1",
+				}: "e2-medium", // for this node type
+			},
+			expected: map[NodeIdentifier]*Node{
+				NodeIdentifier{
+					Cluster:    "cluster1",
+					Name:       "node1",
+					ProviderID: "prov_node1",
+				}: &Node{
+					Cluster:      "cluster1",
+					Name:         "node1",
+					ProviderID:   "prov_node1",
+					NodeType:     "e2-medium",
+					CPUCost:      0.048 * (partialCPUMap["e2-medium"] / 6.0), // adjustmentFactor is (v / GKE cores)
+					CPUCores:     partialCPUMap["e2-medium"],
+					CPUBreakdown: &ClusterCostsBreakdown{},
+					RAMBreakdown: &ClusterCostsBreakdown{},
+				},
+			},
+		},
 	}
 
 	for _, testCase := range cases {