Add external_ip and network flags to configure VM networking in Cloud Image Tests

cmd/manager/main.go

@@ -112,6 +112,9 @@ var (
 	acceleratorType         = flag.String("accelerator_type", "", "Accelerator type to be used for accelerator tests")
 	allImageFamilies        = flag.String("all_image_families", "", "Single image project to test all image families in.")
 	architectureType        = flag.String("architecture_type", "", "Specific architecture to test on. Accepts one of x86 or arm64.")
+	externalIP              = flag.String("external_ip", "", "External IP to use for VMs (ephemeral or none). defaults to ephemeral.")
+	networkFlag             = flag.String("network", "", "The network/VPC to use for VMs. Defaults to default network global/networks/default.")
+	subnet                  = flag.String("subnet", "", "The subnet to use for VMs. Defaults to default network global/networks/default. If the network is in auto subnet mode, the subnetwork is optional. If the network is in custom subnet mode, then this flag should be specified.")
 
 	// zonesRoundRobinIdx points to an index in the list of zones.
 	// This is used to distribute tests across the list of zones in a round robin fashion,
@@ -250,6 +253,11 @@ func main() {
 		return
 	}
 
+	if *externalIP != "" && !strings.EqualFold(*externalIP, "ephemeral") && !strings.EqualFold(*externalIP, "none") {
+		log.Fatal("external_ip must be blank, ephemeral, or none")
+		return
+	}
+
 	var testProjectsReal []string
 	if *testProjects == "" {
 		testProjectsReal = append(testProjectsReal, *project)
@@ -576,6 +584,9 @@ func main() {
 				ReservationURLs:         reservationURLSlice,
 				AcceleratorType:         *acceleratorType,
 				ArgZoneOverride:         *argZoneOverride,
+				ExternalIP:              *externalIP,
+				Network:                 *networkFlag,
+				Subnet:                  *subnet,
 			}, testPackage.setupFunc)
 			if err != nil {
 				log.Fatalf("Failed to create test workflow: %v", err)

testworkflow.go

@@ -104,6 +104,12 @@ type TestWorkflowOpts struct {
 	// true, the zone from the command line will be enforced if the test suite
 	// does specify a zone. If false, the hardcoded zone will be used.
 	ArgZoneOverride bool
+	// ExternalIP is the external IP to use for VMs (ephemeral or none).
+	ExternalIP string
+	// Network is the network to use for VMs. This is used for tests that do not configure custom network interfaces.
+	Network string
+	// Subnet is the subnet to use for VMs. This is used for tests that do not configure custom network interfaces.
+	Subnet string
 }
 
 // TestWorkflow defines a test workflow which creates at least one test VM.
@@ -618,6 +624,91 @@ func (t *TestWorkflow) appendCreateSubnetworksStep(subnetwork *daisy.Subnetwork)
 	return createSubnetworksStep, subnetwork, nil
 }
 
+func (t *TestWorkflow) configureNetwork(step *daisy.Step) {
+	if t.opts == nil {
+		return
+	}
+	if t.opts.Network == "" && t.opts.Subnet == "" {
+		return
+	}
+	if step.CreateInstances == nil {
+		return
+	}
+	for _, instance := range step.CreateInstances.Instances {
+		if instance.Instance.NetworkInterfaces == nil || len(instance.Instance.NetworkInterfaces) == 0 {
+			instance.Instance.NetworkInterfaces = []*compute.NetworkInterface{{}}
+			if t.opts.Network != "" {
+				instance.Instance.NetworkInterfaces[0].Network = t.opts.Network
+			}
+			if t.opts.Subnet != "" {
+				instance.Instance.NetworkInterfaces[0].Subnetwork = t.opts.Subnet
+			}
+		}
+	}
+	for _, instance := range step.CreateInstances.InstancesBeta {
+		if instance.Instance.NetworkInterfaces == nil || len(instance.Instance.NetworkInterfaces) == 0 {
+			instance.Instance.NetworkInterfaces = []*computeBeta.NetworkInterface{{}}
+			if t.opts.Network != "" {
+				instance.Instance.NetworkInterfaces[0].Network = t.opts.Network
+			}
+			if t.opts.Subnet != "" {
+				instance.Instance.NetworkInterfaces[0].Subnetwork = t.opts.Subnet
+			}
+		}
+	}
+}
+
+func (t *TestWorkflow) configureExternalIP(step *daisy.Step) {
+	if t.opts == nil {
+		return
+	}
+	if step.CreateInstances == nil {
+		return
+	}
+
+	if t.opts.ExternalIP == "" {
+		return
+	}
+
+	// Ensure NetworkInterfaces is initialized for all instances.
+	for _, instance := range step.CreateInstances.Instances {
+		if instance.Instance.NetworkInterfaces == nil || len(instance.Instance.NetworkInterfaces) == 0 {
+			instance.Instance.NetworkInterfaces = []*compute.NetworkInterface{{}}
+		}
+	}
+	for _, instance := range step.CreateInstances.InstancesBeta {
+		if instance.Instance.NetworkInterfaces == nil || len(instance.Instance.NetworkInterfaces) == 0 {
+			instance.Instance.NetworkInterfaces = []*computeBeta.NetworkInterface{{}}
+		}
+	}
+
+	if strings.EqualFold(t.opts.ExternalIP, "none") {
+		daisy.UpdateInstanceNoExternalIP(step)
+	} else {
+		// Ephemeral or empty
+		for _, instance := range step.CreateInstances.Instances {
+			for _, networkInterface := range instance.Instance.NetworkInterfaces {
+				if networkInterface.AccessConfigs == nil {
+					networkInterface.AccessConfigs = []*compute.AccessConfig{{Type: "ONE_TO_ONE_NAT"}}
+				}
+				if len(networkInterface.AccessConfigs) > 0 {
+					networkInterface.AccessConfigs[0].NatIP = ""
+				}
+			}
+		}
+		for _, instance := range step.CreateInstances.InstancesBeta {
+			for _, networkInterface := range instance.Instance.NetworkInterfaces {
+				if networkInterface.AccessConfigs == nil {
+					networkInterface.AccessConfigs = []*computeBeta.AccessConfig{{Type: "ONE_TO_ONE_NAT"}}
+				}
+				if len(networkInterface.AccessConfigs) > 0 {
+					networkInterface.AccessConfigs[0].NatIP = ""
+				}
+			}
+		}
+	}
+}
+
 func getGCSPrefix(ctx context.Context, storageClient *storage.Client, project, gcsPath string) (string, error) {
 	// Set global client.
 	client = storageClient
@@ -690,6 +781,13 @@ func finalizeWorkflows(ctx context.Context, tests []*TestWorkflow, gcsPrefix, lo
 			}
 		}
 
+		for _, step := range twf.wf.Steps {
+			if step.CreateInstances != nil {
+				twf.configureNetwork(step)
+				twf.configureExternalIP(step)
+			}
+		}
+
 		// Assume amd64 when arch is not set.
 		arch := "amd64"
 		if twf.Image.Architecture == "ARM64" {

testworkflow_test.go

@@ -15,6 +15,7 @@
 package imagetest
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 	"slices"
@@ -1055,3 +1056,227 @@ func TestRecreateTestWorkflow(t *testing.T) {
 		t.Errorf("recreated name = %q, want %q", recreated.Name, twf.Name)
 	}
 }
+
+func TestFinalizeWorkflowsExternalIP(t *testing.T) {
+	testCases := []struct {
+		name       string
+		externalIP string
+		isBeta     bool
+		validate   func(t *testing.T, step *daisy.Step)
+	}{
+		{
+			name:       "none_v1",
+			externalIP: "none",
+			validate: func(t *testing.T, step *daisy.Step) {
+				inst := step.CreateInstances.Instances[0]
+				if inst.Instance.NetworkInterfaces == nil || len(inst.Instance.NetworkInterfaces) == 0 {
+					t.Fatal("NetworkInterfaces not initialized")
+				}
+				if len(inst.Instance.NetworkInterfaces[0].AccessConfigs) != 0 {
+					t.Errorf("expected 0 AccessConfigs for none, got %d", len(inst.Instance.NetworkInterfaces[0].AccessConfigs))
+				}
+			},
+		},
+		{
+			name:       "none_beta",
+			externalIP: "none",
+			isBeta:     true,
+			validate: func(t *testing.T, step *daisy.Step) {
+				inst := step.CreateInstances.InstancesBeta[0]
+				if inst.Instance.NetworkInterfaces == nil || len(inst.Instance.NetworkInterfaces) == 0 {
+					t.Fatal("NetworkInterfaces not initialized")
+				}
+				if len(inst.Instance.NetworkInterfaces[0].AccessConfigs) != 0 {
+					t.Errorf("expected 0 AccessConfigs for none, got %d", len(inst.Instance.NetworkInterfaces[0].AccessConfigs))
+				}
+			},
+		},
+		{
+			name:       "ephemeral_v1",
+			externalIP: "ephemeral",
+			validate: func(t *testing.T, step *daisy.Step) {
+				inst := step.CreateInstances.Instances[0]
+				if inst.Instance.NetworkInterfaces == nil || len(inst.Instance.NetworkInterfaces) == 0 {
+					t.Fatal("NetworkInterfaces not initialized")
+				}
+				nic := inst.Instance.NetworkInterfaces[0]
+				if len(nic.AccessConfigs) != 1 {
+					t.Fatalf("expected 1 AccessConfig, got %d", len(nic.AccessConfigs))
+				}
+				if nic.AccessConfigs[0].NatIP != "" {
+					t.Errorf("expected empty NatIP for ephemeral, got %q", nic.AccessConfigs[0].NatIP)
+				}
+			},
+		},
+		{
+			name:       "empty_externalIP_v1",
+			externalIP: "",
+			validate: func(t *testing.T, step *daisy.Step) {
+				inst := step.CreateInstances.Instances[0]
+				if inst.Instance.NetworkInterfaces != nil && len(inst.Instance.NetworkInterfaces) > 0 {
+					t.Errorf("expected NetworkInterfaces to be nil or empty for empty externalIP, got %v", inst.Instance.NetworkInterfaces)
+				}
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			twf := NewTestWorkflowForUnitTest(tc.name, "image", "30m")
+			twf.opts = &TestWorkflowOpts{ExternalIP: tc.externalIP}
+
+			step := &daisy.Step{
+				CreateInstances: &daisy.CreateInstances{},
+			}
+
+			if tc.isBeta {
+				step.CreateInstances.InstancesBeta = []*daisy.InstanceBeta{
+					{
+						Instance: computeBeta.Instance{Name: "vm1"},
+					},
+				}
+			} else {
+				step.CreateInstances.Instances = []*daisy.Instance{
+					{
+						Instance: compute.Instance{Name: "vm1"},
+					},
+				}
+			}
+
+			twf.wf.Steps = map[string]*daisy.Step{"create-vms": step}
+
+			err := finalizeWorkflows(context.Background(), []*TestWorkflow{twf}, "gs://bucket", "/tmp")
+			if err != nil {
+				t.Fatalf("finalizeWorkflows failed: %v", err)
+			}
+
+			tc.validate(t, step)
+		})
+	}
+}
+
+func TestFinalizeWorkflowsNetwork(t *testing.T) {
+	testCases := []struct {
+		name             string
+		network          string
+		subnetwork       string
+		isBeta           bool
+		existingNICs     []*compute.NetworkInterface
+		existingNICsBeta []*computeBeta.NetworkInterface
+		validate         func(t *testing.T, step *daisy.Step)
+	}{
+		{
+			name:    "custom_network_v1",
+			network: "global/networks/custom-net",
+			validate: func(t *testing.T, step *daisy.Step) {
+				inst := step.CreateInstances.Instances[0]
+				if len(inst.Instance.NetworkInterfaces) != 1 {
+					t.Fatalf("expected 1 NetworkInterface, got %d", len(inst.Instance.NetworkInterfaces))
+				}
+				if inst.Instance.NetworkInterfaces[0].Network != "global/networks/custom-net" {
+					t.Errorf("expected Network to be %q, got %q", "global/networks/custom-net", inst.Instance.NetworkInterfaces[0].Network)
+				}
+			},
+		},
+		{
+			name:       "custom_subnetwork_v1",
+			subnetwork: "regions/us-central1/subnetworks/custom-subnet",
+			validate: func(t *testing.T, step *daisy.Step) {
+				inst := step.CreateInstances.Instances[0]
+				if len(inst.Instance.NetworkInterfaces) != 1 {
+					t.Fatalf("expected 1 NetworkInterface, got %d", len(inst.Instance.NetworkInterfaces))
+				}
+				if inst.Instance.NetworkInterfaces[0].Subnetwork != "regions/us-central1/subnetworks/custom-subnet" {
+					t.Errorf("expected Subnetwork to be %q, got %q", "regions/us-central1/subnetworks/custom-subnet", inst.Instance.NetworkInterfaces[0].Subnetwork)
+				}
+			},
+		},
+		{
+			name:       "custom_both_v1",
+			network:    "global/networks/custom-net",
+			subnetwork: "regions/us-central1/subnetworks/custom-subnet",
+			validate: func(t *testing.T, step *daisy.Step) {
+				inst := step.CreateInstances.Instances[0]
+				nic := inst.Instance.NetworkInterfaces[0]
+				if nic.Network != "global/networks/custom-net" {
+					t.Errorf("expected Network to be %q, got %q", "global/networks/custom-net", nic.Network)
+				}
+				if nic.Subnetwork != "regions/us-central1/subnetworks/custom-subnet" {
+					t.Errorf("expected Subnetwork to be %q, got %q", "regions/us-central1/subnetworks/custom-subnet", nic.Subnetwork)
+				}
+			},
+		},
+		{
+			name:    "custom_network_beta",
+			network: "global/networks/custom-net",
+			isBeta:  true,
+			validate: func(t *testing.T, step *daisy.Step) {
+				inst := step.CreateInstances.InstancesBeta[0]
+				if len(inst.Instance.NetworkInterfaces) != 1 {
+					t.Fatalf("expected 1 NetworkInterface, got %d", len(inst.Instance.NetworkInterfaces))
+				}
+				if inst.Instance.NetworkInterfaces[0].Network != "global/networks/custom-net" {
+					t.Errorf("expected Network to be %q, got %q", "global/networks/custom-net", inst.Instance.NetworkInterfaces[0].Network)
+				}
+			},
+		},
+		{
+			name:    "pre_configured_nic_v1",
+			network: "global/networks/custom-net",
+			existingNICs: []*compute.NetworkInterface{
+				{Network: "global/networks/existing-net"},
+			},
+			validate: func(t *testing.T, step *daisy.Step) {
+				inst := step.CreateInstances.Instances[0]
+				if len(inst.Instance.NetworkInterfaces) != 1 {
+					t.Fatalf("expected 1 NetworkInterface, got %d", len(inst.Instance.NetworkInterfaces))
+				}
+				// Should NOT be overwritten by opts.Network
+				if inst.Instance.NetworkInterfaces[0].Network != "global/networks/existing-net" {
+					t.Errorf("expected Network to be %q, got %q", "global/networks/existing-net", inst.Instance.NetworkInterfaces[0].Network)
+				}
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			twf := NewTestWorkflowForUnitTest(tc.name, "image", "30m")
+			twf.opts = &TestWorkflowOpts{
+				Network: tc.network,
+				Subnet:  tc.subnetwork,
+			}
+
+			step := &daisy.Step{
+				CreateInstances: &daisy.CreateInstances{},
+			}
+
+			if tc.isBeta {
+				inst := &daisy.InstanceBeta{
+					Instance: computeBeta.Instance{Name: "vm1"},
+				}
+				if tc.existingNICsBeta != nil {
+					inst.Instance.NetworkInterfaces = tc.existingNICsBeta
+				}
+				step.CreateInstances.InstancesBeta = []*daisy.InstanceBeta{inst}
+			} else {
+				inst := &daisy.Instance{
+					Instance: compute.Instance{Name: "vm1"},
+				}
+				if tc.existingNICs != nil {
+					inst.Instance.NetworkInterfaces = tc.existingNICs
+				}
+				step.CreateInstances.Instances = []*daisy.Instance{inst}
+			}
+
+			twf.wf.Steps = map[string]*daisy.Step{"create-vms": step}
+
+			err := finalizeWorkflows(context.Background(), []*TestWorkflow{twf}, "gs://bucket", "/tmp")
+			if err != nil {
+				t.Fatalf("finalizeWorkflows failed: %v", err)
+			}
+
+			tc.validate(t, step)
+		})
+	}
+}