Repository: kubernetes-sigs/cloud-provider-azure
Pull request #592 (2021-04-15 14:47:22)

feat: add update vm interface
<!--  Thanks for sending a pull request!  Here are some tips for you:

1. If this is your first time, please read our contributor guidelines: https://git.k8s.io/community/contributors/guide#your-first-contribution and developer guide https://git.k8s.io/community/contributors/devel/development.md#development-guide
2. Please label this pull request according to what type of issue you are addressing, especially if this is a release targeted pull request. For reference on required PR/issue labels, read here:
https://git.k8s.io/community/contributors/devel/release.md#issue-kind-label
3. Ensure you have added or ran the appropriate tests for your PR: https://git.k8s.io/community/contributors/devel/testing.md
4. If you want *faster* PR reviews, read how: https://git.k8s.io/community/contributors/guide/pull-requests.md#best-practices-for-faster-reviews
5. Follow the instructions for writing a release note: https://git.k8s.io/community/contributors/guide/release-notes.md
6. If the PR is unfinished, see how to mark it: https://git.k8s.io/community/contributors/guide/pull-requests.md#marking-unfinished-pull-requests
-->

**What type of PR is this?**
/kind feature

**What this PR does / why we need it**:
feat: add update vm interface

**Which issue(s) this PR fixes**:
<!-- 
*Automatically closes linked issue when PR is merged.
Usage: `Fixes #<issue number>`, or `Fixes (paste link of issue)`.
_If PR is about `failing-tests or flakes`, please post the related issues/tests in a comment and do not use `Fixes`_*
-->
Fixes #

**Special notes for your reviewer**:
/hold

**Release note**:
```
feat: add update vm interface
```


k8s-ci-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/kubernetes-sigs/cloud-provider-azure/pull/592#" title="Author self-approved">andyzhangx</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=kubernetes-sigs%2Fcloud-provider-azure).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/kubernetes-sigs/cloud-provider-azure/blob/master/OWNERS)~~ [andyzhangx]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->

Patch:
diff --git a/pkg/provider/azure_controller_common.go b/pkg/provider/azure_controller_common.go
index 4961660a3..12a9e8ef5 100644
--- a/pkg/provider/azure_controller_common.go
+++ b/pkg/provider/azure_controller_common.go
@@ -362,6 +362,18 @@ func (c *controllerCommon) DetachDisk(diskName, diskURI string, nodeName types.N
 	return nil
 }
 
+// UpdateVM updates a vm
+func (c *controllerCommon) UpdateVM(nodeName types.NodeName) error {
+	vmset, err := c.getNodeVMSet(nodeName, azcache.CacheReadTypeUnsafe)
+	if err != nil {
+		return err
+	}
+	node := strings.ToLower(string(nodeName))
+	c.lockMap.LockEntry(node)
+	defer c.lockMap.UnlockEntry(node)
+	return vmset.UpdateVM(nodeName)
+}
+
 func (c *controllerCommon) insertDetachDiskRequest(diskName, diskURI, nodeName string) error {
 	var diskMap map[string]string
 	detachDiskMapKey := nodeName + detachDiskMapKeySuffix
@@ -407,23 +419,23 @@ func (c *controllerCommon) cleanDetachDiskRequests(nodeName string) (map[string]
 }
 
 // getNodeDataDisks invokes vmSet interfaces to get data disks for the node.
-func (c *controllerCommon) getNodeDataDisks(nodeName types.NodeName, crt azcache.AzureCacheReadType) ([]compute.DataDisk, error) {
+func (c *controllerCommon) getNodeDataDisks(nodeName types.NodeName, crt azcache.AzureCacheReadType) ([]compute.DataDisk, *string, error) {
 	vmset, err := c.getNodeVMSet(nodeName, crt)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	return vmset.GetDataDisks(nodeName, crt)
 }
 
 // GetDiskLun finds the lun on the host that the vhd is attached to, given a vhd's diskName and diskURI.
-func (c *controllerCommon) GetDiskLun(diskName, diskURI string, nodeName types.NodeName) (int32, error) {
+func (c *controllerCommon) GetDiskLun(diskName, diskURI string, nodeName types.NodeName) (int32, *string, error) {
 	// getNodeDataDisks need to fetch the cached data/fresh data if cache expired here
 	// to ensure we get LUN based on latest entry.
-	disks, err := c.getNodeDataDisks(nodeName, azcache.CacheReadTypeDefault)
+	disks, provisioningState, err := c.getNodeDataDisks(nodeName, azcache.CacheReadTypeDefault)
 	if err != nil {
 		klog.Errorf("error of getting data disks for node %q: %v", nodeName, err)
-		return -1, err
+		return -1, provisioningState, err
 	}
 
 	for _, disk := range disks {
@@ -435,17 +447,17 @@ func (c *controllerCommon) GetDiskLun(diskName, diskURI string, nodeName types.N
 			} else {
 				// found the disk
 				klog.V(2).Infof("azureDisk - find disk: lun %d name %q uri %q", *disk.Lun, diskName, diskURI)
-				return *disk.Lun, nil
+				return *disk.Lun, provisioningState, nil
 			}
 		}
 	}
-	return -1, fmt.Errorf("cannot find Lun for disk %s", diskName)
+	return -1, provisioningState, fmt.Errorf("cannot find Lun for disk %s", diskName)
 }
 
 // SetDiskLun find unused luns and allocate lun for every disk in diskMap.
 // Return lun of diskURI, -1 if all luns are used.
 func (c *controllerCommon) SetDiskLun(nodeName types.NodeName, diskURI string, diskMap map[string]*AttachDiskOptions) (int32, error) {
-	disks, err := c.getNodeDataDisks(nodeName, azcache.CacheReadTypeForceRefresh)
+	disks, _, err := c.getNodeDataDisks(nodeName, azcache.CacheReadTypeForceRefresh)
 	if err != nil {
 		klog.Errorf("error of getting data disks for node %q: %v", nodeName, err)
 		return -1, err
@@ -519,7 +531,7 @@ func (c *controllerCommon) DisksAreAttached(diskNames []string, nodeName types.N
 	// for every reconcile call. The cache is invalidated after Attach/Detach
 	// disk. So the new entry will be fetched and cached the first time reconcile
 	// loop runs after the Attach/Disk OP which will reflect the latest model.
-	disks, err := c.getNodeDataDisks(nodeName, azcache.CacheReadTypeUnsafe)
+	disks, _, err := c.getNodeDataDisks(nodeName, azcache.CacheReadTypeUnsafe)
 	if err != nil {
 		if errors.Is(err, cloudprovider.InstanceNotFound) {
 			// if host doesn't exist, no need to detach
diff --git a/pkg/provider/azure_controller_standard.go b/pkg/provider/azure_controller_standard.go
index 4d2a072eb..1c203e944 100644
--- a/pkg/provider/azure_controller_standard.go
+++ b/pkg/provider/azure_controller_standard.go
@@ -207,16 +207,40 @@ func (as *availabilitySet) DetachDisk(nodeName types.NodeName, diskMap map[strin
 	return nil
 }
 
+// UpdateVM updates a vm
+func (as *availabilitySet) UpdateVM(nodeName types.NodeName) error {
+	vmName := mapNodeNameToVMName(nodeName)
+	nodeResourceGroup, err := as.GetNodeResourceGroup(vmName)
+	if err != nil {
+		return err
+	}
+	klog.V(2).Infof("azureDisk - update(%s): vm(%s)", nodeResourceGroup, vmName)
+	ctx, cancel := getContextWithCancel()
+	defer cancel()
+
+	// Invalidate the cache right after updating
+	defer func() {
+		_ = as.cloud.vmCache.Delete(vmName)
+	}()
+
+	rerr := as.VirtualMachinesClient.Update(ctx, nodeResourceGroup, vmName, compute.VirtualMachineUpdate{}, "update_vm")
+	klog.V(2).Infof("azureDisk - update(%s): vm(%s) - returned with %v", nodeResourceGroup, vmName, rerr)
+	if rerr != nil {
+		return rerr.Error()
+	}
+	return nil
+}
+
 // GetDataDisks gets a list of data disks attached to the node.
-func (as *availabilitySet) GetDataDisks(nodeName types.NodeName, crt azcache.AzureCacheReadType) ([]compute.DataDisk, error) {
+func (as *availabilitySet) GetDataDisks(nodeName types.NodeName, crt azcache.AzureCacheReadType) ([]compute.DataDisk, *string, error) {
 	vm, err := as.getVirtualMachine(nodeName, crt)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	if vm.StorageProfile.DataDisks == nil {
-		return nil, nil
+		return nil, nil, nil
 	}
 
-	return *vm.StorageProfile.DataDisks, nil
+	return *vm.StorageProfile.DataDisks, vm.ProvisioningState, nil
 }
diff --git a/pkg/provider/azure_controller_vmss.go b/pkg/provider/azure_controller_vmss.go
index d60d3830f..afed55f47 100644
--- a/pkg/provider/azure_controller_vmss.go
+++ b/pkg/provider/azure_controller_vmss.go
@@ -212,16 +212,47 @@ func (ss *ScaleSet) DetachDisk(nodeName types.NodeName, diskMap map[string]strin
 	return nil
 }
 
+// UpdateVM updates a vm
+func (ss *ScaleSet) UpdateVM(nodeName types.NodeName) error {
+	vmName := mapNodeNameToVMName(nodeName)
+	ssName, instanceID, _, err := ss.getVmssVM(vmName, azcache.CacheReadTypeDefault)
+	if err != nil {
+		return err
+	}
+
+	nodeResourceGroup, err := ss.GetNodeResourceGroup(vmName)
+	if err != nil {
+		return err
+	}
+
+	ctx, cancel := getContextWithCancel()
+	defer cancel()
+
+	// Invalidate the cache right after updating
+	defer func() {
+		_ = ss.deleteCacheForNode(vmName)
+	}()
+
+	klog.V(2).Infof("azureDisk - update(%s): vm(%s)", nodeResourceGroup, nodeName)
+	rerr := ss.VirtualMachineScaleSetVMsClient.Update(ctx, nodeResourceGroup, ssName, instanceID, compute.VirtualMachineScaleSetVM{}, "update_vmss_instance")
+
+	klog.V(2).Infof("azureDisk - update(%s): vm(%s) - returned with %v", nodeResourceGroup, nodeName, rerr)
+	if rerr != nil {
+		return rerr.Error()
+	}
+	return nil
+}
+
 // GetDataDisks gets a list of data disks attached to the node.
-func (ss *ScaleSet) GetDataDisks(nodeName types.NodeName, crt azcache.AzureCacheReadType) ([]compute.DataDisk, error) {
+func (ss *ScaleSet) GetDataDisks(nodeName types.NodeName, crt azcache.AzureCacheReadType) ([]compute.DataDisk, *string, error) {
 	_, _, vm, err := ss.getVmssVM(string(nodeName), crt)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	if vm.StorageProfile == nil || vm.StorageProfile.DataDisks == nil {
-		return nil, nil
+		return nil, nil, nil
 	}
 
-	return *vm.StorageProfile.DataDisks, nil
+	return *vm.StorageProfile.DataDisks, vm.ProvisioningState, nil
 }
diff --git a/pkg/provider/azure_mock_vmsets.go b/pkg/provider/azure_mock_vmsets.go
index 975238447..0950d17c0 100644
--- a/pkg/provider/azure_mock_vmsets.go
+++ b/pkg/provider/azure_mock_vmsets.go
@@ -261,12 +261,13 @@ func (mr *MockVMSetMockRecorder) DetachDisk(nodeName, diskMap interface{}) *gomo
 }
 
 // GetDataDisks mocks base method
-func (m *MockVMSet) GetDataDisks(nodeName types.NodeName, string cache.AzureCacheReadType) ([]compute.DataDisk, error) {
+func (m *MockVMSet) GetDataDisks(nodeName types.NodeName, cacheType cache.AzureCacheReadType) ([]compute.DataDisk, *string, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "GetDataDisks", nodeName, string)
+	ret := m.ctrl.Call(m, "GetDataDisks", nodeName, cacheType)
 	ret0, _ := ret[0].([]compute.DataDisk)
-	ret1, _ := ret[1].(error)
-	return ret0, ret1
+	ret1, _ := ret[1].(*string)
+	ret2, _ := ret[2].(error)
+	return ret0, ret1, ret2
 }
 
 // GetDataDisks indicates an expected call of GetDataDisks
@@ -275,6 +276,20 @@ func (mr *MockVMSetMockRecorder) GetDataDisks(nodeName, string interface{}) *gom
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetDataDisks", reflect.TypeOf((*MockVMSet)(nil).GetDataDisks), nodeName, string)
 }
 
+// UpdateVM mocks base method
+func (m *MockVMSet) UpdateVM(nodeName types.NodeName) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "UpdateVM", nodeName)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// UpdateVM indicates an expected call of UpdateVM
+func (mr *MockVMSetMockRecorder) UpdateVM(nodeName interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateVM", reflect.TypeOf((*MockVMSet)(nil).UpdateVM), nodeName)
+}
+
 // GetPowerStatusByNodeName mocks base method
 func (m *MockVMSet) GetPowerStatusByNodeName(name string) (string, error) {
 	m.ctrl.T.Helper()
diff --git a/pkg/provider/azure_vmsets.go b/pkg/provider/azure_vmsets.go
index 0044f64cf..5171c2e6a 100644
--- a/pkg/provider/azure_vmsets.go
+++ b/pkg/provider/azure_vmsets.go
@@ -72,7 +72,10 @@ type VMSet interface {
 	// DetachDisk detaches a disk from vm
 	DetachDisk(nodeName types.NodeName, diskMap map[string]string) error
 	// GetDataDisks gets a list of data disks attached to the node.
-	GetDataDisks(nodeName types.NodeName, string azcache.AzureCacheReadType) ([]compute.DataDisk, error)
+	GetDataDisks(nodeName types.NodeName, string azcache.AzureCacheReadType) ([]compute.DataDisk, *string, error)
+
+	// UpdateVM updates a vm
+	UpdateVM(nodeName types.NodeName) error
 
 	// GetPowerStatusByNodeName returns the power state of the specified node.
 	GetPowerStatusByNodeName(name string) (string, error)


Test patch:
diff --git a/pkg/provider/azure_controller_common_test.go b/pkg/provider/azure_controller_common_test.go
index 6bd14f208..7344a340d 100644
--- a/pkg/provider/azure_controller_common_test.go
+++ b/pkg/provider/azure_controller_common_test.go
@@ -350,6 +350,77 @@ func TestCommonDetachDisk(t *testing.T) {
 	}
 }
 
+func TestCommonUpdateVM(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	testCases := []struct {
+		desc             string
+		vmList           map[string]string
+		nodeName         types.NodeName
+		diskName         string
+		isErrorRetriable bool
+		expectedErr      bool
+	}{
+		{
+			desc:        "error should not be returned if there's no such instance corresponding to given nodeName",
+			nodeName:    "vm1",
+			expectedErr: false,
+		},
+		{
+			desc:             "an error should be returned if vmset detach failed with isErrorRetriable error",
+			vmList:           map[string]string{"vm1": "PowerState/Running"},
+			nodeName:         "vm1",
+			isErrorRetriable: true,
+			expectedErr:      true,
+		},
+		{
+			desc:        "no error shall be returned if there's no matching disk according to given diskName",
+			vmList:      map[string]string{"vm1": "PowerState/Running"},
+			nodeName:    "vm1",
+			diskName:    "disk2",
+			expectedErr: false,
+		},
+		{
+			desc:        "no error shall be returned if the disk exists",
+			vmList:      map[string]string{"vm1": "PowerState/Running"},
+			nodeName:    "vm1",
+			diskName:    "disk1",
+			expectedErr: false,
+		},
+	}
+
+	for i, test := range testCases {
+		testCloud := GetTestCloud(ctrl)
+		common := &controllerCommon{
+			location:              testCloud.Location,
+			storageEndpointSuffix: testCloud.Environment.StorageEndpointSuffix,
+			resourceGroup:         testCloud.ResourceGroup,
+			subscriptionID:        testCloud.SubscriptionID,
+			cloud:                 testCloud,
+			lockMap:               newLockMap(),
+		}
+		expectedVMs := setTestVirtualMachines(testCloud, test.vmList, false)
+		mockVMsClient := testCloud.VirtualMachinesClient.(*mockvmclient.MockInterface)
+		for _, vm := range expectedVMs {
+			mockVMsClient.EXPECT().Get(gomock.Any(), testCloud.ResourceGroup, *vm.Name, gomock.Any()).Return(vm, nil).AnyTimes()
+		}
+		if len(expectedVMs) == 0 {
+			mockVMsClient.EXPECT().Get(gomock.Any(), testCloud.ResourceGroup, gomock.Any(), gomock.Any()).Return(compute.VirtualMachine{}, &retry.Error{HTTPStatusCode: http.StatusNotFound, RawError: cloudprovider.InstanceNotFound}).AnyTimes()
+		}
+		if test.isErrorRetriable {
+			testCloud.CloudProviderBackoff = true
+			testCloud.ResourceRequestBackoff = wait.Backoff{Steps: 1}
+			mockVMsClient.EXPECT().Update(gomock.Any(), testCloud.ResourceGroup, gomock.Any(), gomock.Any(), gomock.Any()).Return(&retry.Error{HTTPStatusCode: http.StatusBadRequest, Retriable: true, RawError: fmt.Errorf("Retriable: true")}).AnyTimes()
+		} else {
+			mockVMsClient.EXPECT().Update(gomock.Any(), testCloud.ResourceGroup, gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+		}
+
+		err := common.UpdateVM(test.nodeName)
+		assert.Equal(t, test.expectedErr, err != nil, "TestCase[%d]: %s, err: %v", i, test.desc, err)
+	}
+}
+
 func TestGetDiskLun(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()
@@ -391,7 +462,7 @@ func TestGetDiskLun(t *testing.T) {
 			mockVMsClient.EXPECT().Get(gomock.Any(), testCloud.ResourceGroup, *vm.Name, gomock.Any()).Return(vm, nil).AnyTimes()
 		}
 
-		lun, err := common.GetDiskLun(test.diskName, test.diskURI, "vm1")
+		lun, _, err := common.GetDiskLun(test.diskName, test.diskURI, "vm1")
 		assert.Equal(t, test.expectedLun, lun, "TestCase[%d]: %s", i, test.desc)
 		assert.Equal(t, test.expectedErr, err != nil, "TestCase[%d]: %s", i, test.desc)
 	}
diff --git a/pkg/provider/azure_controller_standard_test.go b/pkg/provider/azure_controller_standard_test.go
index 1949927c8..254a738fd 100644
--- a/pkg/provider/azure_controller_standard_test.go
+++ b/pkg/provider/azure_controller_standard_test.go
@@ -176,7 +176,67 @@ func TestStandardDetachDisk(t *testing.T) {
 		err := vmSet.DetachDisk(test.nodeName, diskMap)
 		assert.Equal(t, test.expectedError, err != nil, "TestCase[%d]: %s", i, test.desc)
 		if !test.expectedError && test.diskName != "" {
-			dataDisks, err := vmSet.GetDataDisks(test.nodeName, azcache.CacheReadTypeDefault)
+			dataDisks, _, err := vmSet.GetDataDisks(test.nodeName, azcache.CacheReadTypeDefault)
+			assert.Equal(t, true, len(dataDisks) == 1, "TestCase[%d]: %s, err: %v", i, test.desc, err)
+		}
+	}
+}
+
+func TestStandardUpdateVM(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	testCases := []struct {
+		desc          string
+		nodeName      types.NodeName
+		diskName      string
+		isDetachFail  bool
+		expectedError bool
+	}{
+		{
+			desc:          "no error shall be returned if there's no corresponding vm",
+			nodeName:      "vm2",
+			expectedError: false,
+		},
+		{
+			desc:          "no error shall be returned if there's no corresponding disk",
+			nodeName:      "vm1",
+			diskName:      "disk2",
+			expectedError: false,
+		},
+		{
+			desc:          "no error shall be returned if there's a corresponding disk",
+			nodeName:      "vm1",
+			diskName:      "disk1",
+			expectedError: false,
+		},
+		{
+			desc:          "an error shall be returned if detach disk failed",
+			nodeName:      "vm1",
+			isDetachFail:  true,
+			expectedError: true,
+		},
+	}
+
+	for i, test := range testCases {
+		testCloud := GetTestCloud(ctrl)
+		vmSet := testCloud.VMSet
+		expectedVMs := setTestVirtualMachines(testCloud, map[string]string{"vm1": "PowerState/Running"}, false)
+		mockVMsClient := testCloud.VirtualMachinesClient.(*mockvmclient.MockInterface)
+		for _, vm := range expectedVMs {
+			mockVMsClient.EXPECT().Get(gomock.Any(), testCloud.ResourceGroup, *vm.Name, gomock.Any()).Return(vm, nil).AnyTimes()
+		}
+		mockVMsClient.EXPECT().Get(gomock.Any(), testCloud.ResourceGroup, "vm2", gomock.Any()).Return(compute.VirtualMachine{}, &retry.Error{HTTPStatusCode: http.StatusNotFound, RawError: cloudprovider.InstanceNotFound}).AnyTimes()
+		if test.isDetachFail {
+			mockVMsClient.EXPECT().Update(gomock.Any(), testCloud.ResourceGroup, gomock.Any(), gomock.Any(), gomock.Any()).Return(&retry.Error{HTTPStatusCode: http.StatusNotFound, RawError: cloudprovider.InstanceNotFound}).AnyTimes()
+		} else {
+			mockVMsClient.EXPECT().Update(gomock.Any(), testCloud.ResourceGroup, gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+		}
+
+		err := vmSet.UpdateVM(test.nodeName)
+		assert.Equal(t, test.expectedError, err != nil, "TestCase[%d]: %s", i, test.desc)
+		if !test.expectedError && test.diskName != "" {
+			dataDisks, _, err := vmSet.GetDataDisks(test.nodeName, azcache.CacheReadTypeDefault)
 			assert.Equal(t, true, len(dataDisks) == 1, "TestCase[%d]: %s, err: %v", i, test.desc, err)
 		}
 	}
@@ -248,13 +308,13 @@ func TestGetDataDisks(t *testing.T) {
 		mockVMsClient.EXPECT().Get(gomock.Any(), testCloud.ResourceGroup, gomock.Not("vm1"), gomock.Any()).Return(compute.VirtualMachine{}, &retry.Error{HTTPStatusCode: http.StatusNotFound, RawError: cloudprovider.InstanceNotFound}).AnyTimes()
 		mockVMsClient.EXPECT().Update(gomock.Any(), testCloud.ResourceGroup, gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 
-		dataDisks, err := vmSet.GetDataDisks(test.nodeName, test.crt)
+		dataDisks, _, err := vmSet.GetDataDisks(test.nodeName, test.crt)
 		assert.Equal(t, test.expectedDataDisks, dataDisks, "TestCase[%d]: %s", i, test.desc)
 		assert.Equal(t, test.expectedError, err != nil, "TestCase[%d]: %s", i, test.desc)
 
 		if test.crt == azcache.CacheReadTypeUnsafe {
 			time.Sleep(fakeCacheTTL)
-			dataDisks, err := vmSet.GetDataDisks(test.nodeName, test.crt)
+			dataDisks, _, err := vmSet.GetDataDisks(test.nodeName, test.crt)
 			assert.Equal(t, test.expectedDataDisks, dataDisks, "TestCase[%d]: %s", i, test.desc)
 			assert.Equal(t, test.expectedError, err != nil, "TestCase[%d]: %s", i, test.desc)
 		}
diff --git a/pkg/provider/azure_controller_vmss_test.go b/pkg/provider/azure_controller_vmss_test.go
index bb81cd057..6b39142f8 100644
--- a/pkg/provider/azure_controller_vmss_test.go
+++ b/pkg/provider/azure_controller_vmss_test.go
@@ -248,7 +248,110 @@ func TestDetachDiskWithVMSS(t *testing.T) {
 		}
 
 		if !test.expectedErr {
-			dataDisks, err := ss.GetDataDisks(test.vmssvmName, azcache.CacheReadTypeDefault)
+			dataDisks, _, err := ss.GetDataDisks(test.vmssvmName, azcache.CacheReadTypeDefault)
+			assert.Equal(t, true, len(dataDisks) == 1, "TestCase[%d]: %s, actual data disk num: %d, err: %v", i, test.desc, len(dataDisks), err)
+		}
+	}
+}
+
+func TestUpdateVMWithVMSS(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	fakeStatusNotFoundVMSSName := types.NodeName("FakeStatusNotFoundVMSSName")
+	diskName := "disk-name"
+	testCases := []struct {
+		desc           string
+		vmList         map[string]string
+		vmssVMList     []string
+		vmssName       types.NodeName
+		vmssvmName     types.NodeName
+		existedDisk    compute.Disk
+		expectedErr    bool
+		expectedErrMsg error
+	}{
+		{
+			desc:           "an error shall be returned if it is invalid vmss name",
+			vmssVMList:     []string{"vmss-vm-000001"},
+			vmssName:       "vm1",
+			vmssvmName:     "vm1",
+			existedDisk:    compute.Disk{Name: to.StringPtr(diskName)},
+			expectedErr:    true,
+			expectedErrMsg: fmt.Errorf("not a vmss instance"),
+		},
+		{
+			desc:        "no error shall be returned if everything is good",
+			vmssVMList:  []string{"vmss00-vm-000000", "vmss00-vm-000001", "vmss00-vm-000002"},
+			vmssName:    "vmss00",
+			vmssvmName:  "vmss00-vm-000000",
+			existedDisk: compute.Disk{Name: to.StringPtr(diskName)},
+			expectedErr: false,
+		},
+		{
+			desc:           "an error shall be returned if response StatusNotFound",
+			vmssVMList:     []string{"vmss00-vm-000000", "vmss00-vm-000001", "vmss00-vm-000002"},
+			vmssName:       fakeStatusNotFoundVMSSName,
+			vmssvmName:     "vmss00-vm-000000",
+			existedDisk:    compute.Disk{Name: to.StringPtr(diskName)},
+			expectedErr:    true,
+			expectedErrMsg: fmt.Errorf("Retriable: false, RetryAfter: 0s, HTTPStatusCode: 404, RawError: %w", fmt.Errorf("instance not found")),
+		},
+		{
+			desc:        "no error shall be returned if everything is good and the attaching disk does not match data disk",
+			vmssVMList:  []string{"vmss00-vm-000000", "vmss00-vm-000001", "vmss00-vm-000002"},
+			vmssName:    "vmss00",
+			vmssvmName:  "vmss00-vm-000000",
+			existedDisk: compute.Disk{Name: to.StringPtr("disk-name-err")},
+			expectedErr: false,
+		},
+	}
+
+	for i, test := range testCases {
+		scaleSetName := string(test.vmssName)
+		ss, err := NewTestScaleSet(ctrl)
+		assert.NoError(t, err, test.desc)
+		testCloud := ss.cloud
+		testCloud.PrimaryScaleSetName = scaleSetName
+		expectedVMSS := buildTestVMSSWithLB(scaleSetName, "vmss00-vm-", []string{testLBBackendpoolID0}, false)
+		mockVMSSClient := testCloud.VirtualMachineScaleSetsClient.(*mockvmssclient.MockInterface)
+		mockVMSSClient.EXPECT().List(gomock.Any(), testCloud.ResourceGroup).Return([]compute.VirtualMachineScaleSet{expectedVMSS}, nil).AnyTimes()
+		mockVMSSClient.EXPECT().Get(gomock.Any(), testCloud.ResourceGroup, scaleSetName).Return(expectedVMSS, nil).MaxTimes(1)
+		mockVMSSClient.EXPECT().CreateOrUpdate(gomock.Any(), testCloud.ResourceGroup, scaleSetName, gomock.Any()).Return(nil).MaxTimes(1)
+
+		expectedVMSSVMs, _, _ := buildTestVirtualMachineEnv(testCloud, scaleSetName, "", 0, test.vmssVMList, "succeeded", false)
+		for _, vmssvm := range expectedVMSSVMs {
+			vmssvm.StorageProfile = &compute.StorageProfile{
+				OsDisk: &compute.OSDisk{
+					Name: to.StringPtr("osdisk1"),
+					ManagedDisk: &compute.ManagedDiskParameters{
+						ID: to.StringPtr("ManagedID"),
+						DiskEncryptionSet: &compute.DiskEncryptionSetParameters{
+							ID: to.StringPtr("DiskEncryptionSetID"),
+						},
+					},
+				},
+				DataDisks: &[]compute.DataDisk{{
+					Lun:  to.Int32Ptr(0),
+					Name: to.StringPtr(diskName),
+				}},
+			}
+		}
+		mockVMSSVMClient := testCloud.VirtualMachineScaleSetVMsClient.(*mockvmssvmclient.MockInterface)
+		mockVMSSVMClient.EXPECT().List(gomock.Any(), testCloud.ResourceGroup, scaleSetName, gomock.Any()).Return(expectedVMSSVMs, nil).AnyTimes()
+		if scaleSetName == string(fakeStatusNotFoundVMSSName) {
+			mockVMSSVMClient.EXPECT().Update(gomock.Any(), testCloud.ResourceGroup, scaleSetName, gomock.Any(), gomock.Any(), gomock.Any()).Return(&retry.Error{HTTPStatusCode: http.StatusNotFound, RawError: cloudprovider.InstanceNotFound}).AnyTimes()
+		} else {
+			mockVMSSVMClient.EXPECT().Update(gomock.Any(), testCloud.ResourceGroup, scaleSetName, gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+		}
+
+		err = ss.UpdateVM(test.vmssvmName)
+		assert.Equal(t, test.expectedErr, err != nil, "TestCase[%d]: %s, err: %v", i, test.desc, err)
+		if test.expectedErr {
+			assert.EqualError(t, test.expectedErrMsg, err.Error(), "TestCase[%d]: %s, expected error: %v, return error: %v", i, test.desc, test.expectedErrMsg, err)
+		}
+
+		if !test.expectedErr {
+			dataDisks, _, err := ss.GetDataDisks(test.vmssvmName, azcache.CacheReadTypeDefault)
 			assert.Equal(t, true, len(dataDisks) == 1, "TestCase[%d]: %s, actual data disk num: %d, err: %v", i, test.desc, len(dataDisks), err)
 		}
 	}
@@ -334,7 +437,7 @@ func TestGetDataDisksWithVMSS(t *testing.T) {
 		mockVMSSVMClient := testCloud.VirtualMachineScaleSetVMsClient.(*mockvmssvmclient.MockInterface)
 		mockVMSSVMClient.EXPECT().List(gomock.Any(), testCloud.ResourceGroup, scaleSetName, gomock.Any()).Return(expectedVMSSVMs, nil).AnyTimes()
 		mockVMSSVMClient.EXPECT().Update(gomock.Any(), testCloud.ResourceGroup, scaleSetName, gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
-		dataDisks, err := ss.GetDataDisks(test.nodeName, test.crt)
+		dataDisks, _, err := ss.GetDataDisks(test.nodeName, test.crt)
 		assert.Equal(t, test.expectedDataDisks, dataDisks, "TestCase[%d]: %s", i, test.desc)
 		assert.Equal(t, test.expectedErr, err != nil, "TestCase[%d]: %s", i, test.desc)
 		assert.Equal(t, test.expectedErrMsg, err, "TestCase[%d]: %s, expected error: %v, return error: %v", i, test.desc, test.expectedErrMsg, err)