fix scaleQuantity function and tidy up unit tests

Signed-off-by: Noor Malik <[email protected]>

provisioner/template.go

@@ -37,8 +37,6 @@ const (
 	labelsConfigItem = "labels"
 	taintsConfigItem = "taints"
 	dedicatedLabel   = "dedicated"
-
-	giga = 1024 * 1024 * 1024
 )
 
 type templateContext struct {
@@ -769,9 +767,9 @@ func instanceTypeMemory(context *templateContext, instanceType string) (string,
 		return "", err
 	}
 
-	memory := fmt.Sprintf("%v%s", instanceTypeInfo.Memory/giga, "Gi")
+	memory := k8sresource.NewQuantity(instanceTypeInfo.Memory, k8sresource.BinarySI)
 
-	return memory, nil
+	return memory.String(), nil
 }
 
 // scaleQuantity scales a k8sresource.Quantity by a factor, represented as string
@@ -787,8 +785,13 @@ func scaleQuantity(quantityStr string, factor float32) (string, error) {
 	if err != nil {
 		return quantityStr, fmt.Errorf("failed to parse %v as k8sresource.Quantity: %v", quantityStr, err)
 	}
-	// scale the quantity in milli-units to handle fractions
-	quantity.SetMilli(int64(float32(quantity.MilliValue()) * factor))
-
+	// CPU quantities can be represented as milli-cores, so we need to scale them differently
+	if quantity.Value() < 1024 && quantity.Format == k8sresource.DecimalSI {
+		quantity.SetMilli(int64(float32(quantity.MilliValue()) * factor))
+	} else {
+		// memory quantities can be represented as binary SI, so we need to scale them differently
+		scaledValue := int64(float32(quantity.Value()) * factor)
+		quantity.Set(scaledValue)
+	}
 	return quantity.String(), nil
 }

provisioner/template_test.go

@@ -1306,8 +1306,14 @@ func TestScaleQuantity(t *testing.T) {
 }
 
 func TestScaleQuantityError(t *testing.T) {
+	// factor must be positive and non-zero
 	_, err := scaleQuantity("1.0", -1.0)
-	require.Error(t, err)
+	require.Errorf(t, err, "scaling factor must be greater than 0.0")
+
+	// must be a valid k8s quantity like "100m" or "1Gi"
+	quantityStr := "InvalidQuantity"
+	_, err = scaleQuantity(quantityStr, 1.0)
+	require.Errorf(t, err, "failed to parse %v as k8sresource.Quantity: %v", quantityStr, err)
 }
 
 func TestInstanceTypeMemory(t *testing.T) {
@@ -1386,21 +1392,29 @@ func TestScalingTemplate(t *testing.T) {
 		expected string
 	}{
 		{
-			name:     "m5.xlarge",
-			input:    `{{ scaleQuantity (instanceTypeMemory .Values.data.instance_type) 0.1 }}`,
-			data:     map[string]string{"instance_type": "m5.xlarge"},
-			expected: "1.6Gi",
-		},
-		{
-			name:     "c5d.xlarge",
-			input:    `{{ scaleQuantity (instanceTypeMemory .Values.data.instance_type) 0.1 }}`,
-			data:     map[string]string{"instance_type": "c5d.xlarge"},
-			expected: "800Mi",
+			name:  "m5.xlarge",
+			input: `{{ scaleQuantity (instanceTypeMemory .Values.data.instance_type) 0.1 }}`,
+			data:  map[string]string{"instance_type": "m5.xlarge"},
+			// 1.6Gi
+			expected: "1717986944",
 		}, {
-			name:     "m6g.xlarge",
-			input:    `{{ scaleQuantity (instanceTypeMemory .Values.data.instance_type) 0.1 }}`,
-			data:     map[string]string{"instance_type": "m6g.xlarge"},
-			expected: "1.6Gi",
+			name:  "c5d.xlarge",
+			input: `{{ scaleQuantity (instanceTypeMemory .Values.data.instance_type) 0.3 }}`,
+			data:  map[string]string{"instance_type": "c5d.xlarge"},
+			// 2.4Gi
+			expected: "2576980480",
+		}, {
+			name:  "scale CPU evenly",
+			input: `{{ scaleQuantity (instanceTypeCPU .Values.data.instance_type) 0.1 }}`,
+			data:  map[string]string{"instance_type": "m6g.xlarge"},
+			// This is interpreted as 4 * 0.1 = 0.4 => 400m, this is the preferred format for CPU
+			expected: "400m",
+		}, {
+			name:  "scale CPU unevenly",
+			input: `{{ scaleQuantity (instanceTypeCPU .Values.data.instance_type) 0.3 }}`,
+			data:  map[string]string{"instance_type": "m6g.xlarge"},
+			// This is intepreted as 4 * 0.3 = 1.2 => 1200m and not 4/3 = 1.33 => 1333m
+			expected: "1200m",
 		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {