test: refactor unit tests to use common OS VMSS builder

k8s-infra-cherrypick-robot · Jun 3, 2024 · 56e2bd1 · 56e2bd1
1 parent 1fefee3
commit 56e2bd1
Showing 1 changed file with 42 additions and 45 deletions.
diff --git a/pkg/provider/azure_vmss_test.go b/pkg/provider/azure_vmss_test.go
@@ -58,6 +58,42 @@ const (
 	testLBBackendpoolID2   = "/subscriptions/sub/resourceGroups/rg1/providers/Microsoft.Network/loadBalancers/lb/backendAddressPools/backendpool-2"
 )
 
+// helper enum for setting the OS variant
+// of the VMSS image ref.
+type osVersion int
+
+const (
+	unspecified osVersion = iota
+	windows2019
+	windows2022
+	ubuntu
+)
+
+func buildTestOSSpecificVMSSWithLB(name, namePrefix string, lbBackendpoolIDs []string, os osVersion, ipv6 bool) compute.VirtualMachineScaleSet {
+	vmss := buildTestVMSSWithLB(name, namePrefix, lbBackendpoolIDs, ipv6)
+	switch os {
+	case windows2019:
+		vmss.VirtualMachineScaleSetProperties.VirtualMachineProfile.StorageProfile = &compute.VirtualMachineScaleSetStorageProfile{
+			OsDisk: &compute.VirtualMachineScaleSetOSDisk{
+				OsType: compute.OperatingSystemTypesWindows,
+			},
+			ImageReference: &compute.ImageReference{
+				ID: ptr.To("/subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/AKS-Windows/providers/Microsoft.Compute/galleries/AKSWindows/images/windows-2019-containerd/versions/17763.5820.240516"),
+			},
+		}
+	case windows2022:
+		vmss.VirtualMachineScaleSetProperties.VirtualMachineProfile.StorageProfile = &compute.VirtualMachineScaleSetStorageProfile{
+			OsDisk: &compute.VirtualMachineScaleSetOSDisk{
+				OsType: compute.OperatingSystemTypesWindows,
+			},
+			ImageReference: &compute.ImageReference{
+				ID: ptr.To("/subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/AKS-Windows/providers/Microsoft.Compute/galleries/AKSWindows/images/windows-2022-containerd/versions/20348.5820.240516"),
+			},
+		}
+	}
+	return vmss
+}
+
 func buildTestVMSSWithLB(name, namePrefix string, lbBackendpoolIDs []string, ipv6 bool) compute.VirtualMachineScaleSet {
 	lbBackendpoolsV4, lbBackendpoolsV6 := make([]compute.SubResource, 0), make([]compute.SubResource, 0)
 	for _, id := range lbBackendpoolIDs {
@@ -2335,7 +2371,7 @@ func TestEnsureVMSSInPool(t *testing.T) {
 		vmSetName             string
 		clusterIP             string
 		nodes                 []*v1.Node
-		mutate                func(compute.VirtualMachineScaleSet)
+		os                    osVersion
 		isBasicLB             bool
 		isVMSSDeallocating    bool
 		isVMSSNilNICConfig    bool
@@ -2453,16 +2489,7 @@ func TestEnsureVMSSInPool(t *testing.T) {
 			expectedPutVMSS: false,
 			setIPv6Config:   false,
 			expectedErr:     nil,
-			mutate: func(vmss compute.VirtualMachineScaleSet) {
-				vmss.VirtualMachineScaleSetProperties.VirtualMachineProfile.StorageProfile = &compute.VirtualMachineScaleSetStorageProfile{
-					OsDisk: &compute.VirtualMachineScaleSetOSDisk{
-						OsType: compute.OperatingSystemTypesWindows,
-					},
-					ImageReference: &compute.ImageReference{
-						ID: ptr.To("/subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/AKS-Windows/providers/Microsoft.Compute/galleries/AKSWindows/images/windows-2019-containerd/versions/17763.5820.240516"),
-					},
-				}
-			},
+			os:              windows2019,
 		},
 		{
 			description: "ensureVMSSInPool should add Windows2019 VM to IPv4 backend pool even if service is IPv6",
@@ -2479,16 +2506,7 @@ func TestEnsureVMSSInPool(t *testing.T) {
 			expectedPutVMSS: true,
 			setIPv6Config:   false,
 			expectedErr:     nil,
-			mutate: func(vmss compute.VirtualMachineScaleSet) {
-				vmss.VirtualMachineScaleSetProperties.VirtualMachineProfile.StorageProfile = &compute.VirtualMachineScaleSetStorageProfile{
-					OsDisk: &compute.VirtualMachineScaleSetOSDisk{
-						OsType: compute.OperatingSystemTypesWindows,
-					},
-					ImageReference: &compute.ImageReference{
-						ID: ptr.To("/subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/AKS-Windows/providers/Microsoft.Compute/galleries/AKSWindows/images/windows-2019-containerd/versions/17763.5820.240516"),
-					},
-				}
-			},
+			os:              windows2019,
 		},
 		{
 			description: "ensureVMSSInPool should add Windows 2022 VM to IPv6 backend pool",
@@ -2505,16 +2523,7 @@ func TestEnsureVMSSInPool(t *testing.T) {
 			expectedPutVMSS: true,
 			setIPv6Config:   true,
 			expectedErr:     nil,
-			mutate: func(vmss compute.VirtualMachineScaleSet) {
-				vmss.VirtualMachineScaleSetProperties.VirtualMachineProfile.StorageProfile = &compute.VirtualMachineScaleSetStorageProfile{
-					OsDisk: &compute.VirtualMachineScaleSetOSDisk{
-						OsType: compute.OperatingSystemTypesWindows,
-					},
-					ImageReference: &compute.ImageReference{
-						ID: ptr.To("/subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/AKS-Windows/providers/Microsoft.Compute/galleries/AKSWindows/images/windows-2022-containerd/versions/20348.5820.240516"),
-					},
-				}
-			},
+			os:              windows2022,
 		},
 		{
 			description: "ensureVMSSInPool should fail if no IPv6 network config - Windows 2022",
@@ -2531,16 +2540,7 @@ func TestEnsureVMSSInPool(t *testing.T) {
 			expectedPutVMSS: false,
 			setIPv6Config:   false,
 			expectedErr:     fmt.Errorf("failed to find a primary IP configuration (IPv6=true) for the VMSS VM or VMSS \"vmss\""),
-			mutate: func(vmss compute.VirtualMachineScaleSet) {
-				vmss.VirtualMachineScaleSetProperties.VirtualMachineProfile.StorageProfile = &compute.VirtualMachineScaleSetStorageProfile{
-					OsDisk: &compute.VirtualMachineScaleSetOSDisk{
-						OsType: compute.OperatingSystemTypesWindows,
-					},
-					ImageReference: &compute.ImageReference{
-						ID: ptr.To("/subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/AKS-Windows/providers/Microsoft.Compute/galleries/AKSWindows/images/windows-2022-containerd/versions/20348.5820.240516"),
-					},
-				}
-			},
+			os:              windows2022,
 		},
 		{
 			description: "ensureVMSSInPool should update the VMSS correctly for IPv6",
@@ -2626,10 +2626,7 @@ func TestEnsureVMSSInPool(t *testing.T) {
 				ss.LoadBalancerSku = consts.LoadBalancerSkuStandard
 			}
 
-			expectedVMSS := buildTestVMSSWithLB(testVMSSName, "vmss-vm-", []string{testLBBackendpoolID0}, test.setIPv6Config)
-			if test.mutate != nil {
-				test.mutate(expectedVMSS)
-			}
+			expectedVMSS := buildTestOSSpecificVMSSWithLB(testVMSSName, "vmss-vm-", []string{testLBBackendpoolID0}, test.os, test.setIPv6Config)
 			if test.isVMSSDeallocating {
 				expectedVMSS.ProvisioningState = pointer.String(consts.ProvisionStateDeleting)
 			}