Refactor and add tests

pkg/cloud/gcpprovider.go

@@ -1395,6 +1395,33 @@ func (gcp *gcpKey) GPUType() string {
 	return ""
 }
 
+func parseGCPInstanceTypeLabel(it string) string {
+	var instanceType string
+
+	splitByDash := strings.Split(it, "-")
+
+	// GKE nodes are labeled with the GCP instance type, but users can deploy on GCP
+	// with tools like K3s, whose instance type labels will be "k3s". This logic
+	// avoids a panic in the slice operation then there are no dashes (-) in the
+	// instance type label value.
+	if len(splitByDash) < 2 {
+		instanceType = "unknown"
+	} else {
+		instanceType = strings.ToLower(strings.Join(splitByDash[:2], ""))
+		if instanceType == "n1highmem" || instanceType == "n1highcpu" {
+			instanceType = "n1standard" // These are priced the same. TODO: support n1ultrahighmem
+		} else if instanceType == "n2highmem" || instanceType == "n2highcpu" {
+			instanceType = "n2standard"
+		} else if instanceType == "e2highmem" || instanceType == "e2highcpu" {
+			instanceType = "e2standard"
+		} else if strings.HasPrefix(instanceType, "custom") {
+			instanceType = "custom" // The suffix of custom does not matter
+		}
+	}
+
+	return instanceType
+}
+
 // GetKey maps node labels to information needed to retrieve pricing data
 func (gcp *gcpKey) Features() string {
 	var instanceType string
@@ -1403,26 +1430,7 @@ func (gcp *gcpKey) Features() string {
 		log.DedupedErrorf(1, "Missing or Unknown 'node.kubernetes.io/instance-type' node label")
 		instanceType = "unknown"
 	} else {
-		splitByDash := strings.Split(it, "-")
-
-		// GKE nodes are labeled with the GCP instance type, but users can deploy on GCP
-		// with tools like K3s, whose instance type labels will be "k3s". This logic
-		// avoids a panic in the slice operation then there are no dashes (-) in the
-		// instance type label value.
-		if len(splitByDash) < 3 {
-			instanceType = "unknown"
-		} else {
-			instanceType = strings.ToLower(strings.Join(splitByDash[:2], ""))
-			if instanceType == "n1highmem" || instanceType == "n1highcpu" {
-				instanceType = "n1standard" // These are priced the same. TODO: support n1ultrahighmem
-			} else if instanceType == "n2highmem" || instanceType == "n2highcpu" {
-				instanceType = "n2standard"
-			} else if instanceType == "e2highmem" || instanceType == "e2highcpu" {
-				instanceType = "e2standard"
-			} else if strings.HasPrefix(instanceType, "custom") {
-				instanceType = "custom" // The suffix of custom does not matter
-			}
-		}
+		instanceType = parseGCPInstanceTypeLabel(it)
 	}
 
 	r, _ := util.GetRegion(gcp.Labels)

pkg/cloud/gcpprovider_test.go

@@ -0,0 +1,36 @@
+package cloud
+
+import (
+	"testing"
+)
+
+func TestParseGCPInstanceTypeLabel(t *testing.T) {
+	cases := []struct {
+		input    string
+		expected string
+	}{
+		{
+			input:    "n1-standard-2",
+			expected: "n1standard",
+		},
+		{
+			input:    "e2-medium",
+			expected: "e2medium",
+		},
+		{
+			input:    "k3s",
+			expected: "unknown",
+		},
+		{
+			input:    "custom-n1-standard-2",
+			expected: "custom",
+		},
+	}
+
+	for _, test := range cases {
+		result := parseGCPInstanceTypeLabel(test.input)
+		if result != test.expected {
+			t.Errorf("Input: %s, Expected: %s, Actual: %s", test.input, test.expected, result)
+		}
+	}
+}