From 7f09dcab0616f70ca96531f1e62c070f2a3872e2 Mon Sep 17 00:00:00 2001
From: Peter Rifel <pgrifel@gmail.com>
Date: Fri, 23 Apr 2021 22:03:28 -0500
Subject: [PATCH] Add more cloud provider flags to e2e.test

---
 tests/e2e/pkg/tester/tester.go | 103 +++++++++++++++++++++++++++++----
 1 file changed, 92 insertions(+), 11 deletions(-)

diff --git a/tests/e2e/pkg/tester/tester.go b/tests/e2e/pkg/tester/tester.go
index 01f3938b23..bc74bf7c05 100644
--- a/tests/e2e/pkg/tester/tester.go
+++ b/tests/e2e/pkg/tester/tester.go
@@ -172,16 +172,93 @@ func (t *Tester) addZoneFlag() error {
 		return nil
 	}
 
+	zoneNames, err := t.getZones()
+	if err != nil {
+		return err
+	}
+
+	// gce-zone only expects one zone, we just pass the first one
+	zone := zoneNames[0]
+	klog.Infof("Setting --gce-zone=%s", zone)
+	t.TestArgs += " --gce-zone=" + zone
+
+	// TODO: Pass the new gce-zones flag for 1.21 with all zones?
+
+	return nil
+}
+
+func (t *Tester) addMultiZoneFlag() error {
+	if hasFlag(t.TestArgs, "gce-multizone") {
+		return nil
+	}
+
+	zoneNames, err := t.getZones()
+	if err != nil {
+		return err
+	}
+
+	klog.Infof("Setting --gce-multizone=%t", len(zoneNames) > 1)
+	t.TestArgs += fmt.Sprintf(" --gce-multizone=%t", len(zoneNames) > 1)
+
+	return nil
+}
+
+func (t *Tester) addRegionFlag() error {
+	// gce-zone is used for other cloud providers as well
+	if hasFlag(t.TestArgs, "gce-region") {
+		return nil
+	}
+
 	cluster, err := t.getKopsCluster()
 	if err != nil {
 		return err
 	}
 
-	igs, err := t.getKopsInstanceGroups()
+	// We don't explicitly set the provider's region in the spec so we need to extract it from vairous fields
+	var region string
+	switch cluster.Spec.CloudProvider {
+	case "aws":
+		zone := cluster.Spec.Subnets[0].Zone
+		region = zone[:len(zone)-1]
+	case "gce":
+		region = cluster.Spec.Subnets[0].Region
+	default:
+		return fmt.Errorf("unhandled region detection for cloud provider: %v", cluster.Spec.CloudProvider)
+	}
+
+	klog.Infof("Setting --gce-region=%s", region)
+	t.TestArgs += " --gce-region=" + region
+	return nil
+}
+
+func (t *Tester) addClusterTagFlag() error {
+	if hasFlag(t.TestArgs, "cluster-tag") {
+		return nil
+	}
+
+	cluster, err := t.getKopsCluster()
 	if err != nil {
 		return err
 	}
 
+	clusterName := cluster.ObjectMeta.Name
+	klog.Infof("Setting --cluster-tag=%s", clusterName)
+	t.TestArgs += " --cluster-tag=" + clusterName
+
+	return nil
+}
+
+func (t *Tester) getZones() ([]string, error) {
+	cluster, err := t.getKopsCluster()
+	if err != nil {
+		return nil, err
+	}
+
+	igs, err := t.getKopsInstanceGroups()
+	if err != nil {
+		return nil, err
+	}
+
 	zones := sets.NewString()
 	// Gather zones on AWS
 	for _, subnet := range cluster.Spec.Subnets {
@@ -198,17 +275,9 @@ func (t *Tester) addZoneFlag() error {
 	zoneNames := zones.List()
 
 	if len(zoneNames) == 0 {
-		return fmt.Errorf("no zones found in instance groups")
+		return nil, fmt.Errorf("no zones found in instance groups")
 	}
-
-	// gce-zone only expects one zone, we just pass the first one
-	zone := zoneNames[0]
-	klog.Infof("Setting --gce-zone=%s", zone)
-	t.TestArgs += " --gce-zone=" + zone
-
-	// TODO: Pass the new gce-zones flag for 1.21 with all zones?
-
-	return nil
+	return zoneNames, nil
 }
 
 func (t *Tester) execute() error {
@@ -244,6 +313,18 @@ func (t *Tester) execute() error {
 		return err
 	}
 
+	if err := t.addClusterTagFlag(); err != nil {
+		return err
+	}
+
+	if err := t.addRegionFlag(); err != nil {
+		return err
+	}
+
+	if err := t.addMultiZoneFlag(); err != nil {
+		return err
+	}
+
 	return t.Test()
 }