Add unit tests for getting vmss node IP

pkg/cloudprovider/providers/azure/azure_fakes.go

@@ -207,6 +207,13 @@ func (fAPC *fakeAzurePIPClient) List(ctx context.Context, resourceGroupName stri
 	return value, nil
 }
 
+func (fAPC *fakeAzurePIPClient) setFakeStore(store map[string]map[string]network.PublicIPAddress) {
+	fAPC.mutex.Lock()
+	defer fAPC.mutex.Unlock()
+
+	fAPC.FakeStore = store
+}
+
 type fakeAzureInterfacesClient struct {
 	mutex     *sync.Mutex
 	FakeStore map[string]map[string]network.Interface
@@ -247,7 +254,24 @@ func (fIC *fakeAzureInterfacesClient) Get(ctx context.Context, resourceGroupName
 }
 
 func (fIC *fakeAzureInterfacesClient) GetVirtualMachineScaleSetNetworkInterface(ctx context.Context, resourceGroupName string, virtualMachineScaleSetName string, virtualmachineIndex string, networkInterfaceName string, expand string) (result network.Interface, err error) {
-	return result, nil
+	fIC.mutex.Lock()
+	defer fIC.mutex.Unlock()
+	if _, ok := fIC.FakeStore[resourceGroupName]; ok {
+		if entity, ok := fIC.FakeStore[resourceGroupName][networkInterfaceName]; ok {
+			return entity, nil
+		}
+	}
+	return result, autorest.DetailedError{
+		StatusCode: http.StatusNotFound,
+		Message:    "Not such Interface",
+	}
+}
+
+func (fIC *fakeAzureInterfacesClient) setFakeStore(store map[string]map[string]network.Interface) {
+	fIC.mutex.Lock()
+	defer fIC.mutex.Unlock()
+
+	fIC.FakeStore = store
 }
 
 type fakeAzureVirtualMachinesClient struct {

pkg/cloudprovider/providers/azure/azure_vmss_test.go

@@ -21,9 +21,16 @@ import (
 	"testing"
 
 	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2018-04-01/compute"
+	"github.com/Azure/azure-sdk-for-go/services/network/mgmt/2017-09-01/network"
+	"github.com/Azure/go-autorest/autorest/to"
 	"github.com/stretchr/testify/assert"
 )
 
+const (
+	fakePrivateIP = "10.240.0.10"
+	fakePublicIP  = "10.10.10.10"
+)
+
 func newTestScaleSet(scaleSetName, zone string, faultDomain int32, vmList []string) (*scaleSet, error) {
 	cloud := getTestCloud()
 	setTestVirtualMachineCloud(cloud, scaleSetName, zone, faultDomain, vmList)
@@ -37,6 +44,11 @@ func newTestScaleSet(scaleSetName, zone string, faultDomain int32, vmList []stri
 
 func setTestVirtualMachineCloud(ss *Cloud, scaleSetName, zone string, faultDomain int32, vmList []string) {
 	virtualMachineScaleSetsClient := newFakeVirtualMachineScaleSetsClient()
+	virtualMachineScaleSetVMsClient := newFakeVirtualMachineScaleSetVMsClient()
+	publicIPAddressesClient := newFakeAzurePIPClient("rg")
+	interfaceClient := newFakeAzureInterfacesClient()
+
+	// set test scale sets.
 	scaleSets := make(map[string]map[string]compute.VirtualMachineScaleSet)
 	scaleSets["rg"] = map[string]compute.VirtualMachineScaleSet{
 		scaleSetName: {
@@ -45,17 +57,27 @@ func setTestVirtualMachineCloud(ss *Cloud, scaleSetName, zone string, faultDomai
 	}
 	virtualMachineScaleSetsClient.setFakeStore(scaleSets)
 
-	virtualMachineScaleSetVMsClient := newFakeVirtualMachineScaleSetVMsClient()
+	testInterfaces := map[string]map[string]network.Interface{
-	ssVMs := make(map[string]map[string]compute.VirtualMachineScaleSetVM)
+		"rg": make(map[string]network.Interface),
-	ssVMs["rg"] = make(map[string]compute.VirtualMachineScaleSetVM)
+	}
+	testPIPs := map[string]map[string]network.PublicIPAddress{
+		"rg": make(map[string]network.PublicIPAddress),
+	}
+	ssVMs := map[string]map[string]compute.VirtualMachineScaleSetVM{
+		"rg": make(map[string]compute.VirtualMachineScaleSetVM),
+	}
 	for i := range vmList {
-		ID := fmt.Sprintf("/subscriptions/script/resourceGroups/rg/providers/Microsoft.Compute/virtualMachineScaleSets/%s/virtualMachines/%d", scaleSetName, i)
 		nodeName := vmList[i]
+		ID := fmt.Sprintf("/subscriptions/script/resourceGroups/rg/providers/Microsoft.Compute/virtualMachineScaleSets/%s/virtualMachines/%d", scaleSetName, i)
+		interfaceID := fmt.Sprintf("/subscriptions/script/resourceGroups/rg/providers/Microsoft.Compute/virtualMachineScaleSets/%s/virtualMachines/%d/networkInterfaces/%s", scaleSetName, i, nodeName)
+		publicAddressID := fmt.Sprintf("/subscriptions/script/resourceGroups/rg/providers/Microsoft.Compute/virtualMachineScaleSets/%s/virtualMachines/%d/networkInterfaces/%s/ipConfigurations/ipconfig1/publicIPAddresses/%s", scaleSetName, i, nodeName, nodeName)
 		instanceID := fmt.Sprintf("%d", i)
 		vmName := fmt.Sprintf("%s_%s", scaleSetName, instanceID)
+
+		// set vmss virtual machine.
 		networkInterfaces := []compute.NetworkInterfaceReference{
 			{
-				ID: &nodeName,
+				ID: &interfaceID,
 			},
 		}
 		vmssVM := compute.VirtualMachineScaleSetVM{
@@ -75,17 +97,46 @@ func setTestVirtualMachineCloud(ss *Cloud, scaleSetName, zone string, faultDomai
 			Name:       &vmName,
 			Location:   &ss.Location,
 		}
-
 		if zone != "" {
 			zones := []string{zone}
 			vmssVM.Zones = &zones
 		}
 		ssVMs["rg"][vmName] = vmssVM
+
+		// set interfaces.
+		testInterfaces["rg"][nodeName] = network.Interface{
+			ID: &interfaceID,
+			InterfacePropertiesFormat: &network.InterfacePropertiesFormat{
+				IPConfigurations: &[]network.InterfaceIPConfiguration{
+					{
+						InterfaceIPConfigurationPropertiesFormat: &network.InterfaceIPConfigurationPropertiesFormat{
+							Primary:          to.BoolPtr(true),
+							PrivateIPAddress: to.StringPtr(fakePrivateIP),
+							PublicIPAddress: &network.PublicIPAddress{
+								ID: to.StringPtr(publicAddressID),
+							},
+						},
+					},
+				},
+			},
+		}
+
+		// set public IPs.
+		testPIPs["rg"][nodeName] = network.PublicIPAddress{
+			ID: to.StringPtr(publicAddressID),
+			PublicIPAddressPropertiesFormat: &network.PublicIPAddressPropertiesFormat{
+				IPAddress: to.StringPtr(fakePublicIP),
+			},
+		}
 	}
 	virtualMachineScaleSetVMsClient.setFakeStore(ssVMs)
+	interfaceClient.setFakeStore(testInterfaces)
+	publicIPAddressesClient.setFakeStore(testPIPs)
 
 	ss.VirtualMachineScaleSetsClient = virtualMachineScaleSetsClient
 	ss.VirtualMachineScaleSetVMsClient = virtualMachineScaleSetVMsClient
+	ss.InterfacesClient = interfaceClient
+	ss.PublicIPAddressesClient = publicIPAddressesClient
 }
 
 func TestGetScaleSetVMInstanceID(t *testing.T) {
@@ -216,3 +267,43 @@ func TestGetZoneByNodeName(t *testing.T) {
 		assert.Equal(t, test.expected, real.FailureDomain, test.description)
 	}
 }
+
+func TestGetIPByNodeName(t *testing.T) {
+	testCases := []struct {
+		description string
+		scaleSet    string
+		vmList      []string
+		nodeName    string
+		expected    []string
+		expectError bool
+	}{
+		{
+			description: "GetIPByNodeName should get node's privateIP and publicIP",
+			scaleSet:    "ss",
+			vmList:      []string{"vmssee6c2000000", "vmssee6c2000001"},
+			nodeName:    "vmssee6c2000000",
+			expected:    []string{fakePrivateIP, fakePublicIP},
+		},
+		{
+			description: "GetIPByNodeName should return error for non-exist nodes",
+			scaleSet:    "ss",
+			vmList:      []string{"vmssee6c2000000", "vmssee6c2000001"},
+			nodeName:    "agente6c2000005",
+			expectError: true,
+		},
+	}
+
+	for _, test := range testCases {
+		ss, err := newTestScaleSet(test.scaleSet, "", 0, test.vmList)
+		assert.NoError(t, err, test.description)
+
+		privateIP, publicIP, err := ss.GetIPByNodeName(test.nodeName)
+		if test.expectError {
+			assert.Error(t, err, test.description)
+			continue
+		}
+
+		assert.NoError(t, err, test.description)
+		assert.Equal(t, test.expected, []string{privateIP, publicIP}, test.description)
+	}
+}