ソースを参照

Applied feedback

kamilkrampa 3 年 前
コミット
92273e8767
2 ファイル変更14 行追加25 行削除
  1. 6 17
      pkg/cloud/gcpprovider.go
  2. 8 8
      pkg/cloud/gcpprovider_test.go

+ 6 - 17
pkg/cloud/gcpprovider.go

@@ -1257,12 +1257,7 @@ func (gcp *gcpKey) ID() string {
 
 func (gcp *gcpKey) GPUType() string {
 	if t, ok := gcp.Labels[GKE_GPU_TAG]; ok {
-		var usageType string
-		if isPreemptible(gcp.Labels) {
-			usageType = "preemptible"
-		} else {
-			usageType = "ondemand"
-		}
+		usageType := getUsageType(gcp.Labels)
 		log.Debugf("GPU of type: \"%s\" found", t)
 		return t + "," + usageType
 	}
@@ -1309,13 +1304,7 @@ func (gcp *gcpKey) Features() string {
 
 	r, _ := util.GetRegion(gcp.Labels)
 	region := strings.ToLower(r)
-	var usageType string
-
-	if isPreemptible(gcp.Labels) {
-		usageType = "preemptible"
-	} else {
-		usageType = "ondemand"
-	}
+	usageType := getUsageType(gcp.Labels)
 
 	if _, ok := gcp.Labels[GKE_GPU_TAG]; ok {
 		return region + "," + instanceType + "," + usageType + "," + "gpu"
@@ -1411,12 +1400,12 @@ func parseGCPProjectID(id string) string {
 	return ""
 }
 
-func isPreemptible(labels map[string]string) bool {
+func getUsageType(labels map[string]string) string {
 	if t, ok := labels["cloud.google.com/gke-preemptible"]; ok && t == "true" {
-		return true
+		return "preemptible"
 	} else if t, ok := labels["cloud.google.com/gke-spot"]; ok && t == "true" {
 		// https://cloud.google.com/kubernetes-engine/docs/concepts/spot-vms
-		return true
+		return "preemptible"
 	}
-	return false
+	return "ondemand"
 }

+ 8 - 8
pkg/cloud/gcpprovider_test.go

@@ -66,39 +66,39 @@ func TestParseGCPProjectID(t *testing.T) {
 	}
 }
 
-func TestIsPreemptible(t *testing.T) {
+func TestGetUsageType(t *testing.T) {
 	cases := []struct {
 		input    map[string]string
-		expected bool
+		expected string
 	}{
 		{
 			input: map[string]string{
 				"cloud.google.com/gke-preemptible": "true",
 			},
-			expected: true,
+			expected: "preemptible",
 		},
 		{
 			input: map[string]string{
 				"cloud.google.com/gke-spot": "true",
 			},
-			expected: true,
+			expected: "preemptible",
 		},
 		{
 			input: map[string]string{
 				"someotherlabel": "true",
 			},
-			expected: false,
+			expected: "ondemand",
 		},
 		{
 			input:    map[string]string{},
-			expected: false,
+			expected: "ondemand",
 		},
 	}
 
 	for _, test := range cases {
-		result := isPreemptible(test.input)
+		result := getUsageType(test.input)
 		if result != test.expected {
-			t.Errorf("Input: %v, Expected: %t, Actual: %t", test.input, test.expected, result)
+			t.Errorf("Input: %v, Expected: %s, Actual: %s", test.input, test.expected, result)
 		}
 	}
 }