Add flags to specify zone/machine-type/OS image

Author: dannawang0221

test/e2e/tests/setup_e2e_test.go

@@ -19,6 +19,7 @@ import (
 	"flag"
 	"fmt"
 	"math/rand"
+	"strings"
 	"testing"
 	"time"
 
@@ -35,6 +36,9 @@ import (
 var (
 	project         = flag.String("project", "", "Project to run tests in")
 	serviceAccount  = flag.String("service-account", "", "Service account to bring up instance with")
+	zones           = flag.String("zones", "us-central1-c,us-central1-b", "Zones to run tests in. If there are multiple zones, separate each by comma")
+	machineType     = flag.String("machine-type", "n1-standard-1", "Type of machine to provision instance on")
+	imageURL        = flag.String("image-url", "projects/debian-cloud/global/images/family/debian-11", "OS image url to get image from")
 	runInProw       = flag.Bool("run-in-prow", false, "If true, use a Boskos loaned project and special CI service accounts and ssh keys")
 	deleteInstances = flag.Bool("delete-instances", false, "Delete the instances after tests run")
 
@@ -60,7 +64,7 @@ var _ = BeforeSuite(func() {
 	tcc := make(chan *remote.TestContext)
 	defer close(tcc)
 
-	zones := []string{"us-central1-c", "us-central1-b"}
+	zones := strings.Split(*zones, ",")
 
 	rand.Seed(time.Now().UnixNano())
 
@@ -89,7 +93,7 @@ var _ = BeforeSuite(func() {
 			nodeID := fmt.Sprintf("gce-pd-csi-e2e-%s", curZone)
 			klog.Infof("Setting up node %s\n", nodeID)
 
-			i, err := remote.SetupInstance(*project, curZone, nodeID, *serviceAccount, computeService)
+			i, err := remote.SetupInstance(*project, curZone, nodeID, *machineType, *serviceAccount, *imageURL, computeService)
 			if err != nil {
 				klog.Fatalf("Failed to setup instance %v: %v", nodeID, err)
 			}

test/remote/instance.go

@@ -38,17 +38,17 @@ import (
 )
 
 const (
-	defaultMachine      = "n1-standard-1"
 	defaultFirewallRule = "default-allow-ssh"
 
 	// timestampFormat is the timestamp format used in the e2e directory name.
 	timestampFormat = "20060102T150405"
 )
 
 type InstanceInfo struct {
-	project string
-	zone    string
-	name    string
+	project     string
+	zone        string
+	name        string
+	machineType string
 
 	// External IP is filled in after instance creation
 	externalIP string
@@ -68,18 +68,19 @@ func (i *InstanceInfo) GetNodeID() string {
 	return common.CreateNodeID(i.project, i.zone, i.name)
 }
 
-func CreateInstanceInfo(project, instanceZone, name string, cs *compute.Service) (*InstanceInfo, error) {
+func CreateInstanceInfo(project, instanceZone, name, machineType string, cs *compute.Service) (*InstanceInfo, error) {
 	return &InstanceInfo{
-		project: project,
-		zone:    instanceZone,
-		name:    name,
+		project:     project,
+		zone:        instanceZone,
+		name:        name,
+		machineType: machineType,
 
 		computeService: cs,
 	}, nil
 }
 
 // Provision a gce instance using image
-func (i *InstanceInfo) CreateOrGetInstance(serviceAccount string) error {
+func (i *InstanceInfo) CreateOrGetInstance(imageURL, serviceAccount string) error {
 	var err error
 	var instance *compute.Instance
 	klog.V(4).Infof("Creating instance: %v", i.name)
@@ -91,10 +92,9 @@ func (i *InstanceInfo) CreateOrGetInstance(serviceAccount string) error {
 		return fmt.Errorf("Failed to create firewall rule: %v", err)
 	}
 
-	imageURL := "projects/debian-cloud/global/images/family/debian-11"
 	inst := &compute.Instance{
 		Name:        i.name,
-		MachineType: machineType(i.zone, ""),
+		MachineType: fmt.Sprintf("zones/%s/machineTypes/%s", i.zone, i.machineType),
 		NetworkInterfaces: []*compute.NetworkInterface{
 			{
 				AccessConfigs: []*compute.AccessConfig{
@@ -215,13 +215,6 @@ func getTimestamp() string {
 	return fmt.Sprintf(time.Now().Format(timestampFormat))
 }
 
-func machineType(zone, machine string) string {
-	if machine == "" {
-		machine = defaultMachine
-	}
-	return fmt.Sprintf("zones/%s/machineTypes/%s", zone, machine)
-}
-
 // Create default SSH filewall rule if it does not exist
 func (i *InstanceInfo) createDefaultFirewallRule() error {
 	var err error

test/remote/setup-teardown.go

@@ -54,14 +54,14 @@ type processes struct {
 }
 
 // SetupInstance sets up the specified GCE Instance for E2E testing and returns a handle to the instance object for future use.
-func SetupInstance(instanceProject, instanceZone, instanceName, instanceServiceAccount string, cs *compute.Service) (*InstanceInfo, error) {
+func SetupInstance(instanceProject, instanceZone, instanceName, instanceMachineType, instanceServiceAccount, instanceImageURL string, cs *compute.Service) (*InstanceInfo, error) {
 	// Create the instance in the requisite zone
-	instance, err := CreateInstanceInfo(instanceProject, instanceZone, instanceName, cs)
+	instance, err := CreateInstanceInfo(instanceProject, instanceZone, instanceName, instanceMachineType, cs)
 	if err != nil {
 		return nil, err
 	}
 
-	err = instance.CreateOrGetInstance(instanceServiceAccount)
+	err = instance.CreateOrGetInstance(instanceImageURL, instanceServiceAccount)
 	if err != nil {
 		return nil, err
 	}