diff --git a/cmd/manager/main.go b/cmd/manager/main.go index 86ab6ca3..f4d3d500 100644 --- a/cmd/manager/main.go +++ b/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) diff --git a/testworkflow.go b/testworkflow.go index 5b7f070d..afb73e6b 100644 --- a/testworkflow.go +++ b/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" { diff --git a/testworkflow_test.go b/testworkflow_test.go index ae2cd1c3..540bf091 100644 --- a/testworkflow_test.go +++ b/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) + }) + } +}