From 28184756c3d8428c686dae1c5908ec263dad52ed Mon Sep 17 00:00:00 2001
From: Justin Santa Barbara <justinsb@google.com>
Date: Mon, 4 Jan 2021 11:26:37 -0500
Subject: [PATCH] kubetest2: add initial support for GCE

Filling in some of the GCE-equivalents to the AWS code.
---
 tests/e2e/kubetest2-kops/deployer/common.go |  8 ++-
 tests/e2e/kubetest2-kops/deployer/up.go     | 28 ++++++--
 tests/e2e/kubetest2-kops/gce/zones.go       | 79 +++++++++++++++++++++
 3 files changed, 107 insertions(+), 8 deletions(-)
 create mode 100644 tests/e2e/kubetest2-kops/gce/zones.go

diff --git a/tests/e2e/kubetest2-kops/deployer/common.go b/tests/e2e/kubetest2-kops/deployer/common.go
index 1a29bb61b7..badb1d036a 100644
--- a/tests/e2e/kubetest2-kops/deployer/common.go
+++ b/tests/e2e/kubetest2-kops/deployer/common.go
@@ -23,6 +23,7 @@ import (
 	"os"
 	"path"
 	"path/filepath"
+	"strings"
 
 	"k8s.io/klog/v2"
 )
@@ -79,6 +80,7 @@ func (d *deployer) verifyKopsFlags() error {
 
 	switch d.CloudProvider {
 	case "aws":
+	case "gce":
 	default:
 		return errors.New("unsupported --cloud-provider value")
 	}
@@ -104,7 +106,11 @@ func (d *deployer) env() []string {
 
 // featureFlags returns the kops feature flags to set
 func (d *deployer) featureFlags() string {
-	return "+SpecOverrideFlag"
+	ff := []string{
+		"+SpecOverrideFlag",
+		"+AlphaAllowGCE",
+	}
+	return strings.Join(ff, ",")
 }
 
 // defaultClusterName returns a kops cluster name to use when ClusterName is not set
diff --git a/tests/e2e/kubetest2-kops/deployer/up.go b/tests/e2e/kubetest2-kops/deployer/up.go
index 33cf1145c7..e34fedca5c 100644
--- a/tests/e2e/kubetest2-kops/deployer/up.go
+++ b/tests/e2e/kubetest2-kops/deployer/up.go
@@ -23,6 +23,7 @@ import (
 
 	"k8s.io/klog/v2"
 	"k8s.io/kops/tests/e2e/kubetest2-kops/aws"
+	"k8s.io/kops/tests/e2e/kubetest2-kops/gce"
 	"k8s.io/kops/tests/e2e/kubetest2-kops/util"
 	"sigs.k8s.io/kubetest2/pkg/exec"
 )
@@ -37,11 +38,6 @@ func (d *deployer) Up() error {
 		return err
 	}
 
-	zones, err := aws.RandomZones(1)
-	if err != nil {
-		return err
-	}
-
 	args := []string{
 		d.KopsBinaryPath, "create", "cluster",
 		"--name", d.ClusterName,
@@ -49,16 +45,34 @@ func (d *deployer) Up() error {
 		"--cloud", d.CloudProvider,
 		"--kubernetes-version", d.KubernetesVersion,
 		"--master-count", "1",
-		"--master-size", "c5.large",
 		"--master-volume-size", "48",
 		"--node-count", "4",
 		"--node-volume-size", "48",
 		"--override", "cluster.spec.nodePortAccess=0.0.0.0/0",
 		"--ssh-public-key", d.SSHPublicKeyPath,
-		"--zones", strings.Join(zones, ","),
 		"--yes",
 	}
 
+	if d.CloudProvider == "aws" {
+		zones, err := aws.RandomZones(1)
+		if err != nil {
+			return err
+		}
+		args = append(args, "--zones", strings.Join(zones, ","))
+
+		args = append(args, "--master-size", "c5.large")
+	}
+
+	if d.CloudProvider == "gce" {
+		zones, err := gce.RandomZones(1)
+		if err != nil {
+			return err
+		}
+		args = append(args, "--zones", strings.Join(zones, ","))
+
+		args = append(args, "--master-size", "e2-standard-2")
+	}
+
 	klog.Info(strings.Join(args, " "))
 	cmd := exec.Command(args[0], args[1:]...)
 	cmd.SetEnv(d.env()...)
diff --git a/tests/e2e/kubetest2-kops/gce/zones.go b/tests/e2e/kubetest2-kops/gce/zones.go
new file mode 100644
index 0000000000..04198db2ac
--- /dev/null
+++ b/tests/e2e/kubetest2-kops/gce/zones.go
@@ -0,0 +1,79 @@
+/*
+Copyright 2020 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package gce
+
+import (
+	"errors"
+	"math/rand"
+	"sort"
+)
+
+var allZones = []string{
+	// Starting with the us- zones (we can add the others later)
+	"us-central1-a",
+	"us-central1-b",
+	"us-central1-c",
+	"us-central1-f",
+	"us-east1-b",
+	"us-east1-c",
+	"us-east1-d",
+	"us-east4-a",
+	"us-east4-b",
+	"us-east4-c",
+	"us-west1-a",
+	"us-west1-b",
+	"us-west1-c",
+	"us-west2-a",
+	"us-west2-b",
+	"us-west2-c",
+	"us-west3-a",
+	"us-west3-b",
+	"us-west3-c",
+	"us-west4-a",
+	"us-west4-b",
+	"us-west4-c",
+}
+
+// ErrNoEligibleRegion indicates the requested number of zones is not available in any region
+var ErrNoEligibleRegion = errors.New("No eligible GCE region found with enough zones")
+
+// RandomZones returns a random set of availability zones within a region
+func RandomZones(count int) ([]string, error) {
+	regions := make(map[string][]string)
+	for _, zone := range allZones {
+		region := zone[:len(zone)-2]
+		regions[region] = append(regions[region], zone)
+	}
+	eligibleRegions := make([][]string, 0)
+	for _, zones := range regions {
+		if len(zones) >= count {
+			eligibleRegions = append(eligibleRegions, zones)
+		}
+	}
+	if len(eligibleRegions) == 0 {
+		return nil, ErrNoEligibleRegion
+	}
+	chosenRegion := eligibleRegions[rand.Int()%len(eligibleRegions)]
+
+	chosenZones := make([]string, 0)
+	randIndexes := rand.Perm(len(chosenRegion))
+	for i := 0; i < count; i++ {
+		chosenZones = append(chosenZones, chosenRegion[randIndexes[i]])
+	}
+	sort.Strings(chosenZones)
+	return chosenZones, nil
+}