Repository: Azure/aad-pod-identity
Pull request #522 (2020-03-04 20:22:25)

switch to using patch instead of createorupdate
<!-- Thank you for helping AAD Pod Identity with a pull request! -->

**Reason for Change**:
<!-- What does this PR improve or fix in AAD Pod Identity? Why is it needed? -->
- Switch to using `PATCH` instead of `CreateOrUpdate`

**Issue Fixed**:
<!-- If this PR fixes GitHub issue 1234, add "Fixes #1234" to the next line. -->
fixes https://github.com/Azure/aad-pod-identity/issues/287

**Notes for Reviewers**:


kkmsft: /azp run
azure-pipelines[bot]: <samp>
Azure Pipelines could not run because the pipeline triggers exclude this branch/path.<br>

</samp>
aramase: /azp run
azure-pipelines[bot]: <samp>
Azure Pipelines could not run because the pipeline triggers exclude this branch/path.<br>

</samp>

Patch:
diff --git a/cmd/demo/main.go b/cmd/demo/main.go
index bef034ca..2b43af1b 100644
--- a/cmd/demo/main.go
+++ b/cmd/demo/main.go
@@ -8,7 +8,8 @@ import (
 	"strings"
 	"time"
 
-	compute "github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2018-04-01/compute"
+	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2019-12-01/compute"
+
 	"k8s.io/klog"
 
 	"github.com/Azure/go-autorest/autorest/adal"
diff --git a/go.mod b/go.mod
index 9fdc8de9..20118005 100644
--- a/go.mod
+++ b/go.mod
@@ -3,11 +3,10 @@ module github.com/Azure/aad-pod-identity
 go 1.12
 
 require (
-	contrib.go.opencensus.io/exporter/ocagent v0.5.0 // indirect
 	contrib.go.opencensus.io/exporter/prometheus v0.1.0
-	github.com/Azure/azure-sdk-for-go v31.0.0+incompatible
-	github.com/Azure/go-autorest/autorest v0.3.0
-	github.com/Azure/go-autorest/autorest/adal v0.1.0
+	github.com/Azure/azure-sdk-for-go v40.4.0+incompatible
+	github.com/Azure/go-autorest/autorest v0.10.0
+	github.com/Azure/go-autorest/autorest/adal v0.8.2
 	github.com/Azure/go-autorest/autorest/azure/auth v0.1.0
 	github.com/Azure/go-autorest/autorest/to v0.2.0 // indirect
 	github.com/Azure/go-autorest/autorest/validation v0.1.0 // indirect
@@ -27,6 +26,7 @@ require (
 	github.com/onsi/gomega v1.4.3
 	github.com/pkg/errors v0.8.0
 	github.com/spf13/pflag v1.0.1
+	github.com/stretchr/testify v1.3.0
 	go.opencensus.io v0.22.0
 	golang.org/x/sync v0.0.0-20190227155943-e225da77a7e6
 	golang.org/x/time v0.0.0-20180412165947-fbb02b2291d2 // indirect
diff --git a/go.sum b/go.sum
index 7e631159..13e72e80 100644
--- a/go.sum
+++ b/go.sum
@@ -1,25 +1,32 @@
 cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw=
 cloud.google.com/go v0.34.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw=
 contrib.go.opencensus.io/exporter/ocagent v0.4.12/go.mod h1:450APlNTSR6FrvC3CTRqYosuDstRB9un7SOx2k/9ckA=
-contrib.go.opencensus.io/exporter/ocagent v0.5.0 h1:TKXjQSRS0/cCDrP7KvkgU6SmILtF/yV2TOs/02K/WZQ=
-contrib.go.opencensus.io/exporter/ocagent v0.5.0/go.mod h1:ImxhfLRpxoYiSq891pBrLVhN+qmP8BTVvdH2YLs7Gl0=
 contrib.go.opencensus.io/exporter/prometheus v0.1.0 h1:SByaIoWwNgMdPSgl5sMqM2KDE5H/ukPWBRo314xiDvg=
 contrib.go.opencensus.io/exporter/prometheus v0.1.0/go.mod h1:cGFniUXGZlKRjzOyuZJ6mgB+PgBcCIa79kEKR8YCW+A=
-github.com/Azure/azure-sdk-for-go v31.0.0+incompatible h1:18nT+M3yxnWcO66yoJyomlCoKMu578UHh0DjJBA5c1M=
-github.com/Azure/azure-sdk-for-go v31.0.0+incompatible/go.mod h1:9XXNKU+eRnpl9moKnB4QOLf1HestfXbmab5FXxiDBjc=
+github.com/Azure/azure-sdk-for-go v40.4.0+incompatible h1:fqIZrZZOfzH4BikHmEtVax+ewacSVCY8u9hXNPUSxF4=
+github.com/Azure/azure-sdk-for-go v40.4.0+incompatible/go.mod h1:9XXNKU+eRnpl9moKnB4QOLf1HestfXbmab5FXxiDBjc=
 github.com/Azure/go-autorest/autorest v0.1.0/go.mod h1:AKyIcETwSUFxIcs/Wnq/C+kwCtlEYGUVd7FPNb2slmg=
-github.com/Azure/go-autorest/autorest v0.3.0 h1:yOmXNB2qa2Kx40wMZB19YyafzjCHacXPk8u0neqa+M0=
-github.com/Azure/go-autorest/autorest v0.3.0/go.mod h1:AKyIcETwSUFxIcs/Wnq/C+kwCtlEYGUVd7FPNb2slmg=
+github.com/Azure/go-autorest/autorest v0.9.0/go.mod h1:xyHB1BMZT0cuDHU7I0+g046+BFDTQ8rEZB0s4Yfa6bI=
+github.com/Azure/go-autorest/autorest v0.10.0 h1:mvdtztBqcL8se7MdrUweNieTNi4kfNG6GOJuurQJpuY=
+github.com/Azure/go-autorest/autorest v0.10.0/go.mod h1:/FALq9T/kS7b5J5qsQ+RSTUdAmGFqi0vUdVNNx8q630=
 github.com/Azure/go-autorest/autorest/adal v0.1.0 h1:RSw/7EAullliqwkZvgIGDYZWQm1PGKXI8c4aY/87yuU=
 github.com/Azure/go-autorest/autorest/adal v0.1.0/go.mod h1:MeS4XhScH55IST095THyTxElntu7WqB7pNbZo8Q5G3E=
+github.com/Azure/go-autorest/autorest/adal v0.5.0/go.mod h1:8Z9fGy2MpX0PvDjB1pEgQTmVqjGhiHBW7RJJEciWzS0=
+github.com/Azure/go-autorest/autorest/adal v0.8.2 h1:O1X4oexUxnZCaEUGsvMnr8ZGj8HI37tNezwY4npRqA0=
+github.com/Azure/go-autorest/autorest/adal v0.8.2/go.mod h1:ZjhuQClTqx435SRJ2iMlOxPYt3d2C/T/7TiQCVZSn3Q=
 github.com/Azure/go-autorest/autorest/azure/auth v0.1.0 h1:YgO/vSnJEc76NLw2ecIXvXa8bDWiqf1pOJzARAoZsYU=
 github.com/Azure/go-autorest/autorest/azure/auth v0.1.0/go.mod h1:Gf7/i2FUpyb/sGBLIFxTBzrNzBo7aPXXE3ZVeDRwdpM=
 github.com/Azure/go-autorest/autorest/azure/cli v0.1.0 h1:YTtBrcb6mhA+PoSW8WxFDoIIyjp13XqJeX80ssQtri4=
 github.com/Azure/go-autorest/autorest/azure/cli v0.1.0/go.mod h1:Dk8CUAt/b/PzkfeRsWzVG9Yj3ps8mS8ECztu43rdU8U=
 github.com/Azure/go-autorest/autorest/date v0.1.0 h1:YGrhWfrgtFs84+h0o46rJrlmsZtyZRg470CqAXTZaGM=
 github.com/Azure/go-autorest/autorest/date v0.1.0/go.mod h1:plvfp3oPSKwf2DNjlBjWF/7vwR+cUD/ELuzDCXwHUVA=
+github.com/Azure/go-autorest/autorest/date v0.2.0 h1:yW+Zlqf26583pE43KhfnhFcdmSWlm5Ew6bxipnr/tbM=
+github.com/Azure/go-autorest/autorest/date v0.2.0/go.mod h1:vcORJHLJEh643/Ioh9+vPmf1Ij9AEBM5FuBIXLmIy0g=
 github.com/Azure/go-autorest/autorest/mocks v0.1.0 h1:Kx+AUU2Te+A3JIyYn6Dfs+cFgx5XorQKuIXrZGoq/SI=
 github.com/Azure/go-autorest/autorest/mocks v0.1.0/go.mod h1:OTyCOPRA2IgIlWxVYxBee2F5Gr4kF2zd2J5cFRaIDN0=
+github.com/Azure/go-autorest/autorest/mocks v0.2.0/go.mod h1:OTyCOPRA2IgIlWxVYxBee2F5Gr4kF2zd2J5cFRaIDN0=
+github.com/Azure/go-autorest/autorest/mocks v0.3.0 h1:qJumjCaCudz+OcqE9/XtEPfvtOjOmKaui4EOpFI6zZc=
+github.com/Azure/go-autorest/autorest/mocks v0.3.0/go.mod h1:a8FDP3DYzQ4RYfVAxAN3SVSiiO77gL2j2ronKKP0syM=
 github.com/Azure/go-autorest/autorest/to v0.2.0 h1:nQOZzFCudTh+TvquAtCRjM01VEYx85e9qbwt5ncW4L8=
 github.com/Azure/go-autorest/autorest/to v0.2.0/go.mod h1:GunWKJp1AEqgMaGLV+iocmRAJWqST1wQYhyyjXJ3SJc=
 github.com/Azure/go-autorest/autorest/validation v0.1.0 h1:ISSNzGUh+ZSzizJWOWzs8bwpXIePbGLW4z/AmUFGH5A=
@@ -28,6 +35,8 @@ github.com/Azure/go-autorest/logger v0.1.0 h1:ruG4BSDXONFRrZZJ2GUXDiUyVpayPmb1Gn
 github.com/Azure/go-autorest/logger v0.1.0/go.mod h1:oExouG+K6PryycPJfVSxi/koC6LSNgds39diKLz7Vrc=
 github.com/Azure/go-autorest/tracing v0.1.0 h1:TRBxC5Pj/fIuh4Qob0ZpkggbfT8RC0SubHbpV3p4/Vc=
 github.com/Azure/go-autorest/tracing v0.1.0/go.mod h1:ROEEAFwXycQw7Sn3DXNtEedEvdeRAgDr0izn4z5Ij88=
+github.com/Azure/go-autorest/tracing v0.5.0 h1:TRn4WjSnkcSy5AEG3pnbtFSwNtwzjr4VYyQflFE619k=
+github.com/Azure/go-autorest/tracing v0.5.0/go.mod h1:r/s2XiOKccPW3HrqB+W0TQzfbtp2fGCgRFtBroKn4Dk=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
 github.com/Shopify/sarama v1.19.0/go.mod h1:FVkBWblsNy7DGZRfXLU0O9RCGt5g3g3yEuWXgklEdEo=
 github.com/Shopify/toxiproxy v2.1.4+incompatible/go.mod h1:OXgGpZ6Cli1/URJOF1DMxUHB2q5Ap20/P/eIdh4G0pI=
@@ -159,6 +168,8 @@ golang.org/x/crypto v0.0.0-20180904163835-0709b304e793/go.mod h1:6SG95UA2DQfeDnf
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20190418165655-df01cb2cc480 h1:O5YqonU5IWby+w98jVUG9h7zlCWCcH4RHyPVReBmhzk=
 golang.org/x/crypto v0.0.0-20190418165655-df01cb2cc480/go.mod h1:WFFai1msRO1wXaEeE5yQxYXgSfI8pQAWXbQop6sCtWE=
+golang.org/x/crypto v0.0.0-20191206172530-e9b2fee46413 h1:ULYEB3JvPRE/IfO+9uO7vKV/xzVTO7XPAwm8xbf4w2g=
+golang.org/x/crypto v0.0.0-20191206172530-e9b2fee46413/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
 golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
 golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE=
 golang.org/x/lint v0.0.0-20190227174305-5b3e6a55c961/go.mod h1:wehouNa3lNwaWXcvxsM5YxQ5yQlVC4a0KAMCusXpPoU=
@@ -174,6 +185,7 @@ golang.org/x/net v0.0.0-20190108225652-1e06a53dbb7e/go.mod h1:mL1N/T3taQHkDXs73r
 golang.org/x/net v0.0.0-20190125091013-d26f9f9a57f3/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
 golang.org/x/net v0.0.0-20190213061140-3a22650c66bd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
 golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
+golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
 golang.org/x/net v0.0.0-20190501004415-9ce7a6920f09 h1:KaQtG+aDELoNmXYas3TVkGNYRuq8JQ1aa7LJt8EXVyo=
 golang.org/x/net v0.0.0-20190501004415-9ce7a6920f09/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
 golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U=
@@ -191,6 +203,7 @@ golang.org/x/sys v0.0.0-20181116152217-5ac8a444bdc5/go.mod h1:STP8DvDyc/dI5b8T5h
 golang.org/x/sys v0.0.0-20181122145206-62eef0e2fa9b/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190403152447-81d4e9dc473e/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20190502145724-3ef323f4f1fd h1:r7DufRZuZbWB7j439YfAzP8RPDa9unLkpwQKUYbIMPI=
 golang.org/x/sys v0.0.0-20190502145724-3ef323f4f1fd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
@@ -206,8 +219,6 @@ golang.org/x/tools v0.0.0-20190226205152-f727befe758c/go.mod h1:9Yl7xja0Znq3iFh3
 golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=
 golang.org/x/tools v0.0.0-20190312170243-e65039ee4138/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=
 google.golang.org/api v0.3.1/go.mod h1:6wY9I6uQWHQ8EM57III9mq/AjF+i8G65rmVagqKMtkk=
-google.golang.org/api v0.4.0 h1:KKgc1aqhV8wDPbDzlDtpvyjZFY3vjz85FP7p4wcQUyI=
-google.golang.org/api v0.4.0/go.mod h1:8k5glujaEP+g9n7WNsDg8QP6cUVNI86fCNMcbazEtwE=
 google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9YwlJXL52JkM=
 google.golang.org/appengine v1.4.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4=
 google.golang.org/genproto v0.0.0-20180817151627-c66870c02cf8/go.mod h1:JiN7NxoALGmiZfu7CAH4rXhgtRTLTxftemlI0sWmxmc=
diff --git a/pkg/cloudprovider/cloudprovider.go b/pkg/cloudprovider/cloudprovider.go
index 17afc24f..070ddd23 100644
--- a/pkg/cloudprovider/cloudprovider.go
+++ b/pkg/cloudprovider/cloudprovider.go
@@ -13,7 +13,7 @@ import (
 	config "github.com/Azure/aad-pod-identity/pkg/config"
 	"github.com/Azure/aad-pod-identity/pkg/utils"
 	"github.com/Azure/aad-pod-identity/version"
-	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2018-04-01/compute"
+	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2019-12-01/compute"
 	"github.com/Azure/go-autorest/autorest"
 	"github.com/Azure/go-autorest/autorest/adal"
 	"github.com/Azure/go-autorest/autorest/azure"
@@ -31,8 +31,6 @@ type Client struct {
 
 // ClientInt client interface
 type ClientInt interface {
-	RemoveUserMSI(userAssignedMSIID, name string, isvmss bool) error
-	AssignUserMSI(userAssignedMSIID, name string, isvmss bool) error
 	UpdateUserMSI(addUserAssignedMSIIDs, removeUserAssignedMSIIDs []string, name string, isvmss bool) error
 	GetUserMSIs(name string, isvmss bool) ([]string, error)
 }
@@ -179,22 +177,17 @@ func (c *Client) UpdateUserMSI(addUserAssignedMSIIDs, removeUserAssignedMSIIDs [
 		info = idH.ResetIdentity()
 	}
 
-	requiresUpdate := false
-	// remove msi ids from the list
-	for _, userAssignedMSIID := range removeUserAssignedMSIIDs {
-		requiresUpdate = true
-		if err := info.RemoveUserIdentity(userAssignedMSIID); err != nil {
-			return fmt.Errorf("could not remove identity from node %s: %v", name, err)
-		}
-	}
+	ids := make(map[string]bool)
 	// add new ids to the list
 	for _, userAssignedMSIID := range addUserAssignedMSIIDs {
-		addedToList := info.AppendUserIdentity(userAssignedMSIID)
-		if !addedToList {
-			klog.V(6).Infof("Identity %s already assigned to node %s. Skipping assignment.", userAssignedMSIID, name)
-		}
-		requiresUpdate = requiresUpdate || addedToList
+		ids[userAssignedMSIID] = true
+	}
+	// remove msi ids from the list
+	for _, userAssignedMSIID := range removeUserAssignedMSIIDs {
+		ids[userAssignedMSIID] = false
 	}
+	requiresUpdate := info.SetUserIdentities(ids)
+
 	if requiresUpdate {
 		klog.Infof("Updating user assigned MSIs on %s", name)
 		timeStarted := time.Now()
@@ -206,62 +199,6 @@ func (c *Client) UpdateUserMSI(addUserAssignedMSIIDs, removeUserAssignedMSIIDs [
 	return nil
 }
 
-//RemoveUserMSI - Use the underlying cloud api calls and remove the given user assigned MSI from the vm.
-func (c *Client) RemoveUserMSI(userAssignedMSIID, name string, isvmss bool) error {
-	idH, updateFunc, err := c.getIdentityResource(name, isvmss)
-	if err != nil {
-		return err
-	}
-
-	info := idH.IdentityInfo()
-	if info == nil {
-		klog.Errorf("Identity null for vm: %s ", name)
-		return fmt.Errorf("identity null for vm: %s ", name)
-	}
-
-	if err := info.RemoveUserIdentity(userAssignedMSIID); err != nil {
-		return fmt.Errorf("could not remove identity from node %s: %v", name, err)
-	}
-
-	if err := updateFunc(); err != nil {
-		klog.Error(err)
-		return err
-	}
-
-	return nil
-}
-
-// AssignUserMSI - Use the underlying cloud api call and add the given user assigned MSI to the vm
-func (c *Client) AssignUserMSI(userAssignedMSIID, name string, isvmss bool) error {
-	// Get the vm using the VmClient
-	// Update the assigned identity into the VM using the CreateOrUpdate
-
-	klog.Infof("Find %s in resource group: %s", name, c.Config.ResourceGroupName)
-	timeStarted := time.Now()
-
-	idH, updateFunc, err := c.getIdentityResource(name, isvmss)
-	if err != nil {
-		return err
-	}
-	klog.V(6).Infof("Get of %s completed in %s", name, time.Since(timeStarted))
-
-	info := idH.IdentityInfo()
-	if info == nil {
-		info = idH.ResetIdentity()
-	}
-
-	if info.AppendUserIdentity(userAssignedMSIID) {
-		timeStarted = time.Now()
-		if err := updateFunc(); err != nil {
-			return err
-		}
-		klog.V(6).Infof("CreateOrUpdate of %s completed in %s", name, time.Since(timeStarted))
-	} else {
-		klog.V(6).Infof("Identity %s already assigned to node %s. Skipping assignment.", userAssignedMSIID, name)
-	}
-	return nil
-}
-
 func (c *Client) getIdentityResource(name string, isvmss bool) (idH IdentityHolder, update func() error, retErr error) {
 	rg := c.Config.ResourceGroupName
 
@@ -272,7 +209,7 @@ func (c *Client) getIdentityResource(name string, isvmss bool) (idH IdentityHold
 		}
 
 		update = func() error {
-			return c.VMSSClient.CreateOrUpdate(rg, name, vmss)
+			return c.VMSSClient.UpdateIdentities(rg, name, vmss)
 		}
 		idH = &vmssIdentityHolder{&vmss}
 		return idH, update, nil
@@ -283,10 +220,9 @@ func (c *Client) getIdentityResource(name string, isvmss bool) (idH IdentityHold
 		return nil, nil, err
 	}
 	update = func() error {
-		return c.VMClient.CreateOrUpdate(rg, name, vm)
+		return c.VMClient.UpdateIdentities(rg, name, vm)
 	}
 	idH = &vmIdentityHolder{&vm}
-
 	return idH, update, nil
 }
 
diff --git a/pkg/cloudprovider/identity.go b/pkg/cloudprovider/identity.go
index 44e98523..f308409d 100644
--- a/pkg/cloudprovider/identity.go
+++ b/pkg/cloudprovider/identity.go
@@ -1,10 +1,7 @@
 package cloudprovider
 
 import (
-	"errors"
 	"strings"
-
-	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2018-04-01/compute"
 )
 
 // IdentityHolder represents a resource that contains an Identity object
@@ -20,77 +17,8 @@ type IdentityHolder interface {
 // have different identity types.
 // This abstracts those differences.
 type IdentityInfo interface {
-	AppendUserIdentity(id string) bool
-	RemoveUserIdentity(id string) error
 	GetUserIdentityList() []string
-}
-
-var (
-	errNotAssigned = errors.New("identity is not assigned to the resource")
-	errNotFound    = errors.New("user assigned identity not found")
-)
-
-// filterUserIdentity provides a common implementation for removing an identity
-// from an identity list.
-func filterUserIdentity(idType *compute.ResourceIdentityType, idList *[]string, id string) error {
-	switch *idType {
-	case compute.ResourceIdentityTypeUserAssigned,
-		compute.ResourceIdentityTypeSystemAssignedUserAssigned:
-	default:
-		return nil
-	}
-
-	origLen := len(*idList)
-	filter(idList, id)
-
-	if len(*idList) >= origLen {
-		return errNotAssigned
-	}
-
-	if len(*idList) != 0 {
-		return nil
-	}
-
-	if *idType == compute.ResourceIdentityTypeSystemAssignedUserAssigned {
-		*idType = compute.ResourceIdentityTypeSystemAssigned
-	} else {
-		*idType = compute.ResourceIdentityTypeNone
-	}
-
-	return nil
-}
-
-func filter(ls *[]string, filter string) {
-	if ls == nil {
-		return
-	}
-
-	for i, v := range *ls {
-		if strings.EqualFold(v, filter) {
-			copy((*ls)[i:], (*ls)[i+1:])
-			*ls = (*ls)[:len(*ls)-1]
-			return
-		}
-	}
-}
-
-// appendUserIdentity provides a common implementation for adding a new identity
-// to an identity object.
-func appendUserIdentity(idType *compute.ResourceIdentityType, idList *[]string, newID string) bool {
-	switch *idType {
-	case compute.ResourceIdentityTypeUserAssigned, compute.ResourceIdentityTypeSystemAssignedUserAssigned:
-		// check if this ID is already in the list
-		if checkIfIDInList(*idList, newID) {
-			return false
-		}
-	case compute.ResourceIdentityTypeSystemAssigned:
-		*idType = compute.ResourceIdentityTypeSystemAssignedUserAssigned
-	default:
-		*idType = compute.ResourceIdentityTypeUserAssigned
-	}
-
-	*idList = append(*idList, newID)
-	return true
+	SetUserIdentities(map[string]bool) bool
 }
 
 func checkIfIDInList(idList []string, desiredID string) bool {
diff --git a/pkg/cloudprovider/vm.go b/pkg/cloudprovider/vm.go
index 53b19f60..b4e2a837 100644
--- a/pkg/cloudprovider/vm.go
+++ b/pkg/cloudprovider/vm.go
@@ -8,7 +8,7 @@ import (
 	"github.com/Azure/aad-pod-identity/pkg/metrics"
 	"github.com/Azure/aad-pod-identity/pkg/stats"
 	"github.com/Azure/aad-pod-identity/version"
-	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2018-04-01/compute"
+	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2019-12-01/compute"
 	"github.com/Azure/go-autorest/autorest"
 	"github.com/Azure/go-autorest/autorest/adal"
 	"github.com/Azure/go-autorest/autorest/azure"
@@ -23,8 +23,8 @@ type VMClient struct {
 
 // VMClientInt is the interface used by "cloudprovider" for interacting with Azure vmas
 type VMClientInt interface {
-	CreateOrUpdate(rg string, nodeName string, vm compute.VirtualMachine) error
 	Get(rgName string, nodeName string) (compute.VirtualMachine, error)
+	UpdateIdentities(rg, nodeName string, vmu compute.VirtualMachine) error
 }
 
 // NewVirtualMachinesClient creates a new vm client.
@@ -53,61 +53,57 @@ func NewVirtualMachinesClient(config config.AzureConfig, spt *adal.ServicePrinci
 	}, nil
 }
 
-// CreateOrUpdate creates a new vm, or if the vm already exists it updates the existing one.
-// This is used by "cloudprovider" to *update* add/remove identities from an already existing vm.
-func (c *VMClient) CreateOrUpdate(rg string, nodeName string, vm compute.VirtualMachine) error {
-	// Set the read-only property of extension to null.
-	vm.Resources = nil
+// Get gets the passed in vm.
+func (c *VMClient) Get(rgName string, nodeName string) (compute.VirtualMachine, error) {
 	ctx := context.Background()
 	begin := time.Now()
 	var err error
 
 	defer func() {
 		if err != nil {
-			c.reporter.ReportCloudProviderOperationError(metrics.PutVMOperationName)
+			c.reporter.ReportCloudProviderOperationError(metrics.GetVMOperationName)
 			return
 		}
-		c.reporter.ReportCloudProviderOperationDuration(metrics.PutVMOperationName, time.Since(begin))
+		c.reporter.ReportCloudProviderOperationDuration(metrics.GetVMOperationName, time.Since(begin))
 	}()
 
-	future, err := c.client.CreateOrUpdate(ctx, rg, nodeName, vm)
-	if err != nil {
-		klog.Error(err)
-		return err
-	}
-
-	err = future.WaitForCompletionRef(ctx, c.client.Client)
+	vm, err := c.client.Get(ctx, rgName, nodeName, "")
 	if err != nil {
 		klog.Error(err)
-		return err
+		return vm, err
 	}
-	stats.UpdateCount(stats.TotalPutCalls, 1)
-	stats.Update(stats.CloudPut, time.Since(begin))
-	return nil
+	stats.UpdateCount(stats.TotalGetCalls, 1)
+	stats.Update(stats.CloudGet, time.Since(begin))
+	return vm, nil
 }
 
-// Get gets the passed in vm.
-func (c *VMClient) Get(rgName string, nodeName string) (compute.VirtualMachine, error) {
+// UpdateIdentities updates the user assigned identities for the provided node
+func (c *VMClient) UpdateIdentities(rg, nodeName string, vm compute.VirtualMachine) error {
+	var future compute.VirtualMachinesUpdateFuture
+	var err error
 	ctx := context.Background()
 	begin := time.Now()
-	var err error
 
 	defer func() {
 		if err != nil {
-			c.reporter.ReportCloudProviderOperationError(metrics.GetVMOperationName)
+			c.reporter.ReportCloudProviderOperationError(metrics.PutVMOperationName)
 			return
 		}
-		c.reporter.ReportCloudProviderOperationDuration(metrics.GetVMOperationName, time.Since(begin))
+		c.reporter.ReportCloudProviderOperationDuration(metrics.PutVMOperationName, time.Since(begin))
 	}()
 
-	vm, err := c.client.Get(ctx, rgName, nodeName, "")
-	if err != nil {
+	if future, err = c.client.Update(ctx, rg, nodeName, compute.VirtualMachineUpdate{
+		Identity: vm.Identity}); err != nil {
+		klog.Errorf("Failed to update VM with error %v", err)
+		return err
+	}
+	if err = future.WaitForCompletionRef(ctx, c.client.Client); err != nil {
 		klog.Error(err)
-		return vm, err
+		return err
 	}
-	stats.UpdateCount(stats.TotalGetCalls, 1)
-	stats.Update(stats.CloudGet, time.Since(begin))
-	return vm, nil
+	stats.UpdateCount(stats.TotalPutCalls, 1)
+	stats.Update(stats.CloudPut, time.Since(begin))
+	return nil
 }
 
 type vmIdentityHolder struct {
@@ -130,33 +126,64 @@ type vmIdentityInfo struct {
 	info *compute.VirtualMachineIdentity
 }
 
-func (i *vmIdentityInfo) RemoveUserIdentity(id string) error {
-	if err := filterUserIdentity(&i.info.Type, i.info.IdentityIds, id); err != nil {
-		return err
-	}
-	// If we have either no identity assigned or have the system assigned identity only, then we need to set the
-	// IdentityIds list as nil.
-	if i.info.Type == compute.ResourceIdentityTypeNone || i.info.Type == compute.ResourceIdentityTypeSystemAssigned {
-		i.info.IdentityIds = nil
+func (i *vmIdentityInfo) GetUserIdentityList() []string {
+	var ids []string
+	if i.info == nil {
+		return ids
 	}
-	// if the identityids is nil and identity type is not set, then set it to ResourceIdentityTypeNone
-	if i.info.IdentityIds == nil && i.info.Type == "" {
-		i.info.Type = compute.ResourceIdentityTypeNone
+	for id := range i.info.UserAssignedIdentities {
+		ids = append(ids, id)
 	}
-	return nil
+	return ids
 }
 
-func (i *vmIdentityInfo) AppendUserIdentity(id string) bool {
-	if i.info.IdentityIds == nil {
-		var ids []string
-		i.info.IdentityIds = &ids
+func (i *vmIdentityInfo) SetUserIdentities(ids map[string]bool) bool {
+	if i.info.UserAssignedIdentities == nil {
+		i.info.UserAssignedIdentities = make(map[string]*compute.VirtualMachineIdentityUserAssignedIdentitiesValue)
 	}
-	return appendUserIdentity(&i.info.Type, i.info.IdentityIds, id)
-}
 
-func (i *vmIdentityInfo) GetUserIdentityList() []string {
-	if i.info.IdentityIds == nil {
-		return []string{}
+	nodeList := make(map[string]bool)
+	// add all current existing ids
+	for id, _ := range i.info.UserAssignedIdentities {
+		nodeList[id] = true
+	}
+
+	// add and remove the new list of identities keeping the same type as before
+	userAssignedIdentities := make(map[string]*compute.VirtualMachineIdentityUserAssignedIdentitiesValue)
+	for id, add := range ids {
+		_, exists := nodeList[id]
+		// already exists on node and want to remove existing identity
+		if exists && !add {
+			userAssignedIdentities[id] = nil
+			delete(nodeList, id)
+		}
+		// doesn't exist on the node and want to add new identity
+		if !exists && add {
+			userAssignedIdentities[id] = &compute.VirtualMachineIdentityUserAssignedIdentitiesValue{}
+			nodeList[id] = true
+		}
+		// exists and add - will already be in the nodeList and no need to patch for it
+		// not exists and delete - no need to patch it as it already doesn't exist
 	}
-	return *i.info.IdentityIds
+
+	// all identities are the node are to be removed
+	if len(nodeList) == 0 {
+		i.info.UserAssignedIdentities = nil
+		if i.info.Type == compute.ResourceIdentityTypeSystemAssignedUserAssigned {
+			i.info.Type = compute.ResourceIdentityTypeSystemAssigned
+		} else {
+			i.info.Type = compute.ResourceIdentityTypeNone
+		}
+		return true
+	}
+
+	if i.info.Type == compute.ResourceIdentityTypeSystemAssigned {
+		i.info.Type = compute.ResourceIdentityTypeSystemAssignedUserAssigned
+	}
+	if i.info.Type == compute.ResourceIdentityTypeNone {
+		i.info.Type = compute.ResourceIdentityTypeUserAssigned
+	}
+
+	i.info.UserAssignedIdentities = userAssignedIdentities
+	return len(i.info.UserAssignedIdentities) > 0
 }
diff --git a/pkg/cloudprovider/vmss.go b/pkg/cloudprovider/vmss.go
index 57a01ff1..a1db164e 100644
--- a/pkg/cloudprovider/vmss.go
+++ b/pkg/cloudprovider/vmss.go
@@ -8,7 +8,7 @@ import (
 	"github.com/Azure/aad-pod-identity/pkg/metrics"
 	"github.com/Azure/aad-pod-identity/pkg/stats"
 	"github.com/Azure/aad-pod-identity/version"
-	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2018-04-01/compute"
+	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2019-12-01/compute"
 	"github.com/Azure/go-autorest/autorest"
 	"github.com/Azure/go-autorest/autorest/adal"
 	"github.com/Azure/go-autorest/autorest/azure"
@@ -23,8 +23,8 @@ type VMSSClient struct {
 
 // VMSSClientInt is the interface used by "cloudprovider" for interacting with Azure vmss
 type VMSSClientInt interface {
-	CreateOrUpdate(rg, name string, vm compute.VirtualMachineScaleSet) error
 	Get(rgName, name string) (compute.VirtualMachineScaleSet, error)
+	UpdateIdentities(rg, vmssName string, vmu compute.VirtualMachineScaleSet) error
 }
 
 // NewVMSSClient creates a new vmss client.
@@ -53,14 +53,12 @@ func NewVMSSClient(config config.AzureConfig, spt *adal.ServicePrincipalToken) (
 	}, nil
 }
 
-// CreateOrUpdate creates a new vmss, or if the vmss already exists it updates the existing one.
-// This is used by "cloudprovider" to *update* add/remove identities from an already existing vmss.
-func (c *VMSSClient) CreateOrUpdate(rg string, vmssName string, vm compute.VirtualMachineScaleSet) error {
-	// Set the read-only property of extension to null.
-	//vm.Resources = nil
+// UpdateIdentities updates the user assigned identities for the provided node
+func (c *VMSSClient) UpdateIdentities(rg, vmssName string, vmssIdentities compute.VirtualMachineScaleSet) error {
+	var future compute.VirtualMachineScaleSetsUpdateFuture
+	var err error
 	ctx := context.Background()
 	begin := time.Now()
-	var err error
 
 	defer func() {
 		if err != nil {
@@ -70,14 +68,12 @@ func (c *VMSSClient) CreateOrUpdate(rg string, vmssName string, vm compute.Virtu
 		c.reporter.ReportCloudProviderOperationDuration(metrics.PutVmssOperationName, time.Since(begin))
 	}()
 
-	future, err := c.client.CreateOrUpdate(ctx, rg, vmssName, vm)
-	if err != nil {
-		klog.Error(err)
+	if future, err = c.client.Update(ctx, rg, vmssName, compute.VirtualMachineScaleSetUpdate{
+		Identity: vmssIdentities.Identity}); err != nil {
+		klog.Errorf("Failed to update VM with error %v", err)
 		return err
 	}
-
-	err = future.WaitForCompletionRef(ctx, c.client.Client)
-	if err != nil {
+	if err = future.WaitForCompletionRef(ctx, c.client.Client); err != nil {
 		klog.Error(err)
 		return err
 	}
@@ -121,7 +117,10 @@ func (h *vmssIdentityHolder) IdentityInfo() IdentityInfo {
 }
 
 func (h *vmssIdentityHolder) ResetIdentity() IdentityInfo {
-	h.vmss.Identity = &compute.VirtualMachineScaleSetIdentity{}
+	h.vmss.Identity = &compute.VirtualMachineScaleSetIdentity{
+		Type:                   compute.ResourceIdentityTypeUserAssigned,
+		UserAssignedIdentities: make(map[string]*compute.VirtualMachineScaleSetIdentityUserAssignedIdentitiesValue),
+	}
 	return h.IdentityInfo()
 }
 
@@ -129,33 +128,64 @@ type vmssIdentityInfo struct {
 	info *compute.VirtualMachineScaleSetIdentity
 }
 
-func (i *vmssIdentityInfo) RemoveUserIdentity(id string) error {
-	if err := filterUserIdentity(&i.info.Type, i.info.IdentityIds, id); err != nil {
-		return err
-	}
-	// If we have either no identity assigned or have the system assigned identity only, then we need to set the
-	// IdentityIds list as nil.
-	if i.info.Type == compute.ResourceIdentityTypeNone || i.info.Type == compute.ResourceIdentityTypeSystemAssigned {
-		i.info.IdentityIds = nil
+func (i *vmssIdentityInfo) GetUserIdentityList() []string {
+	var ids []string
+	if i.info == nil {
+		return ids
 	}
-	// if the identityids is nil and identity type is not set, then set it to ResourceIdentityTypeNone
-	if i.info.IdentityIds == nil && i.info.Type == "" {
-		i.info.Type = compute.ResourceIdentityTypeNone
+	for id := range i.info.UserAssignedIdentities {
+		ids = append(ids, id)
 	}
-	return nil
+	return ids
 }
 
-func (i *vmssIdentityInfo) AppendUserIdentity(id string) bool {
-	if i.info.IdentityIds == nil {
-		var ids []string
-		i.info.IdentityIds = &ids
+func (i *vmssIdentityInfo) SetUserIdentities(ids map[string]bool) bool {
+	if i.info.UserAssignedIdentities == nil {
+		i.info.UserAssignedIdentities = make(map[string]*compute.VirtualMachineScaleSetIdentityUserAssignedIdentitiesValue)
 	}
-	return appendUserIdentity(&i.info.Type, i.info.IdentityIds, id)
-}
 
-func (i *vmssIdentityInfo) GetUserIdentityList() []string {
-	if i.info.IdentityIds == nil {
-		return []string{}
+	nodeList := make(map[string]bool)
+	// add all current existing ids
+	for id, _ := range i.info.UserAssignedIdentities {
+		nodeList[id] = true
+	}
+
+	// add and remove the new list of identities keeping the same type as before
+	userAssignedIdentities := make(map[string]*compute.VirtualMachineScaleSetIdentityUserAssignedIdentitiesValue)
+	for id, add := range ids {
+		_, exists := nodeList[id]
+		// already exists on node and want to remove existing identity
+		if exists && !add {
+			userAssignedIdentities[id] = nil
+			delete(nodeList, id)
+		}
+		// doesn't exist on the node and want to add new identity
+		if !exists && add {
+			userAssignedIdentities[id] = &compute.VirtualMachineScaleSetIdentityUserAssignedIdentitiesValue{}
+			nodeList[id] = true
+		}
+		// exists and add - will already be in the nodeList and no need to patch for it
+		// not exists and delete - no need to patch it as it already doesn't exist
+	}
+
+	// all identities are the node are to be removed
+	if len(nodeList) == 0 {
+		i.info.UserAssignedIdentities = nil
+		if i.info.Type == compute.ResourceIdentityTypeSystemAssignedUserAssigned {
+			i.info.Type = compute.ResourceIdentityTypeSystemAssigned
+		} else {
+			i.info.Type = compute.ResourceIdentityTypeNone
+		}
+		return true
+	}
+
+	if i.info.Type == compute.ResourceIdentityTypeSystemAssigned {
+		i.info.Type = compute.ResourceIdentityTypeSystemAssignedUserAssigned
+	}
+	if i.info.Type == compute.ResourceIdentityTypeNone {
+		i.info.Type = compute.ResourceIdentityTypeUserAssigned
 	}
-	return *i.info.IdentityIds
+
+	i.info.UserAssignedIdentities = userAssignedIdentities
+	return len(i.info.UserAssignedIdentities) > 0
 }
diff --git a/pkg/mic/mic.go b/pkg/mic/mic.go
index b9365588..1feba5ab 100644
--- a/pkg/mic/mic.go
+++ b/pkg/mic/mic.go
@@ -365,7 +365,7 @@ func (c *Client) Sync(exit <-chan struct{}) {
 				bindingsFound = len(*listBindings)
 			}
 			klog.Infof("Work done: %v. Found %d pods, %d ids, %d bindings", workDone, len(listPods), idsFound, bindingsFound)
-			klog.Infof("Total work cycles: %d, out of which work was done in: %d.", totalSyncCycles, totalWorkDoneCycles)
+			klog.Infof("Total work cycles: %d, out of which work was done in: %d", totalSyncCycles, totalWorkDoneCycles)
 			stats.Put(stats.Total, time.Since(begin))
 
 			c.Reporter.Report(


Test patch:
diff --git a/pkg/cloudprovider/cloudprovider_test.go b/pkg/cloudprovider/cloudprovider_test.go
index df1317a1..bde0cef6 100644
--- a/pkg/cloudprovider/cloudprovider_test.go
+++ b/pkg/cloudprovider/cloudprovider_test.go
@@ -1,8 +1,8 @@
 package cloudprovider
 
 import (
-	"flag"
 	"reflect"
+	"sort"
 	"testing"
 
 	"github.com/Azure/aad-pod-identity/pkg/config"
@@ -10,7 +10,7 @@ import (
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
-	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2018-04-01/compute"
+	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2019-12-01/compute"
 )
 
 func TestParseResourceID(t *testing.T) {
@@ -55,10 +55,6 @@ func TestParseResourceID(t *testing.T) {
 	}
 }
 func TestSimple(t *testing.T) {
-	flag.Set("logtostderr", "true")
-	flag.Set("v", "3")
-	flag.Parse()
-
 	vmProvider := "azure:///subscriptions/fakeSub/resourceGroups/fakeGroup/providers/Microsoft.Compute/virtualMachines/node3"
 	vmssProvider := "azure:///subscriptions/fakeSub/resourceGroups/fakeGroup/providers/Microsoft.Compute/virtualMachineScaleSets/node4/virtualMachines/0"
 
@@ -80,13 +76,11 @@ func TestSimple(t *testing.T) {
 			node3 := &corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: "node3-0"}, Spec: corev1.NodeSpec{ProviderID: vmProvider}}
 			node4 := &corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: "node4-vmss0000000"}, Spec: corev1.NodeSpec{ProviderID: vmssProvider}}
 
-			cloudClient.AssignUserMSI("ID0", node0.Name, false)
-			cloudClient.AssignUserMSI("ID0", node0.Name, false)
-			cloudClient.AssignUserMSI("ID0again", node0.Name, false)
-			cloudClient.AssignUserMSI("ID1", node1.Name, false)
-			cloudClient.AssignUserMSI("ID2", node2.Name, false)
-			cloudClient.AssignUserMSI("ID3", node3.Name, false)
-			cloudClient.AssignUserMSI("ID4", node4.Name, true)
+			cloudClient.UpdateUserMSI([]string{"ID0", "ID0again"}, []string{}, node0.Name, false)
+			cloudClient.UpdateUserMSI([]string{"ID1"}, []string{}, node1.Name, false)
+			cloudClient.UpdateUserMSI([]string{"ID2"}, []string{}, node2.Name, false)
+			cloudClient.UpdateUserMSI([]string{"ID3"}, []string{}, node3.Name, false)
+			cloudClient.UpdateUserMSI([]string{"ID4"}, []string{}, node4.Name, true)
 
 			testMSI := []string{"ID0", "ID0again"}
 			if !cloudClient.CompareMSI(node0.Name, false, testMSI) {
@@ -94,8 +88,9 @@ func TestSimple(t *testing.T) {
 				t.Error("MSI mismatch")
 			}
 
-			cloudClient.RemoveUserMSI("ID0", node0.Name, false)
-			cloudClient.RemoveUserMSI("ID2", node2.Name, false)
+			cloudClient.UpdateUserMSI([]string{}, []string{"ID0"}, node0.Name, false)
+			cloudClient.UpdateUserMSI([]string{}, []string{"ID2"}, node2.Name, false)
+
 			testMSI = []string{"ID0again"}
 			if !cloudClient.CompareMSI(node0.Name, false, testMSI) {
 				cloudClient.PrintMSI(t)
@@ -154,6 +149,7 @@ type TestCloudClient struct {
 type TestVMClient struct {
 	*VMClient
 	nodeMap map[string]*compute.VirtualMachine
+	nodeIDs map[string]map[string]bool
 	err     *error
 }
 
@@ -170,15 +166,39 @@ func (c *TestVMClient) Get(rgName string, nodeName string) (ret compute.VirtualM
 	if stored == nil {
 		vm := new(compute.VirtualMachine)
 		c.nodeMap[nodeName] = vm
+		c.nodeIDs[nodeName] = make(map[string]bool)
 		return *vm, nil
 	}
+
+	storedIDs := c.nodeIDs[nodeName]
+	newVMIdentity := make(map[string]*compute.VirtualMachineIdentityUserAssignedIdentitiesValue)
+	for id := range storedIDs {
+		newVMIdentity[id] = &compute.VirtualMachineIdentityUserAssignedIdentitiesValue{}
+	}
+	stored.Identity.UserAssignedIdentities = newVMIdentity
 	return *stored, nil
 }
 
-func (c *TestVMClient) CreateOrUpdate(rg string, nodeName string, vm compute.VirtualMachine) error {
+func (c *TestVMClient) UpdateIdentities(rg, nodeName string, vm compute.VirtualMachine) error {
 	if c.err != nil {
 		return *c.err
 	}
+
+	if vm.Identity != nil && vm.Identity.UserAssignedIdentities != nil {
+		for k, v := range vm.Identity.UserAssignedIdentities {
+			if v == nil {
+				delete(c.nodeIDs[nodeName], k)
+			} else {
+				c.nodeIDs[nodeName][k] = true
+			}
+		}
+	}
+	if vm.Identity != nil && vm.Identity.UserAssignedIdentities == nil {
+		for k := range c.nodeIDs[nodeName] {
+			delete(c.nodeIDs[nodeName], k)
+		}
+	}
+
 	c.nodeMap[nodeName] = &vm
 	return nil
 }
@@ -187,7 +207,11 @@ func (c *TestVMClient) ListMSI() (ret map[string]*[]string) {
 	ret = make(map[string]*[]string)
 
 	for key, val := range c.nodeMap {
-		ret[key] = val.Identity.IdentityIds
+		var ids []string
+		for k := range val.Identity.UserAssignedIdentities {
+			ids = append(ids, k)
+		}
+		ret[key] = &ids
 	}
 	return ret
 }
@@ -198,19 +222,26 @@ func (c *TestVMClient) CompareMSI(nodeName string, userIDs []string) bool {
 		return false
 	}
 
-	ids := stored.Identity.IdentityIds
+	var ids []string
+	for k := range c.nodeIDs[nodeName] {
+		ids = append(ids, k)
+	}
 	if ids == nil {
 		if len(userIDs) == 0 && stored.Identity.Type == compute.ResourceIdentityTypeNone { // Validate that we have reset the resource type as none.
 			return true
 		}
 		return false
 	}
-	return reflect.DeepEqual(*ids, userIDs)
+
+	sort.Strings(ids)
+	sort.Strings(userIDs)
+	return reflect.DeepEqual(ids, userIDs)
 }
 
 type TestVMSSClient struct {
 	*VMSSClient
 	nodeMap map[string]*compute.VirtualMachineScaleSet
+	nodeIDs map[string]map[string]bool
 	err     *error
 }
 
@@ -227,16 +258,39 @@ func (c *TestVMSSClient) Get(rgName string, nodeName string) (ret compute.Virtua
 	if stored == nil {
 		vm := new(compute.VirtualMachineScaleSet)
 		c.nodeMap[nodeName] = vm
+		c.nodeIDs[nodeName] = make(map[string]bool)
 		return *vm, nil
 	}
+
+	storedIDs := c.nodeIDs[nodeName]
+	newVMSSIdentity := make(map[string]*compute.VirtualMachineScaleSetIdentityUserAssignedIdentitiesValue)
+	for id := range storedIDs {
+		newVMSSIdentity[id] = &compute.VirtualMachineScaleSetIdentityUserAssignedIdentitiesValue{}
+	}
+	stored.Identity.UserAssignedIdentities = newVMSSIdentity
 	return *stored, nil
 }
 
-func (c *TestVMSSClient) CreateOrUpdate(rg string, nodeName string, vm compute.VirtualMachineScaleSet) error {
+func (c *TestVMSSClient) UpdateIdentities(rg, nodeName string, vmss compute.VirtualMachineScaleSet) error {
 	if c.err != nil {
 		return *c.err
 	}
-	c.nodeMap[nodeName] = &vm
+	if vmss.Identity != nil && vmss.Identity.UserAssignedIdentities != nil {
+		for k, v := range vmss.Identity.UserAssignedIdentities {
+			if v == nil {
+				delete(c.nodeIDs[nodeName], k)
+			} else {
+				c.nodeIDs[nodeName][k] = true
+			}
+		}
+	}
+	if vmss.Identity != nil && vmss.Identity.UserAssignedIdentities == nil {
+		for k := range c.nodeIDs[nodeName] {
+			delete(c.nodeIDs[nodeName], k)
+		}
+	}
+
+	c.nodeMap[nodeName] = &vmss
 	return nil
 }
 
@@ -244,7 +298,11 @@ func (c *TestVMSSClient) ListMSI() (ret map[string]*[]string) {
 	ret = make(map[string]*[]string)
 
 	for key, val := range c.nodeMap {
-		ret[key] = val.Identity.IdentityIds
+		var ids []string
+		for k := range val.Identity.UserAssignedIdentities {
+			ids = append(ids, k)
+		}
+		ret[key] = &ids
 	}
 	return ret
 }
@@ -255,14 +313,21 @@ func (c *TestVMSSClient) CompareMSI(nodeName string, userIDs []string) bool {
 		return false
 	}
 
-	ids := stored.Identity.IdentityIds
+	var ids []string
+	for k := range c.nodeIDs[nodeName] {
+		ids = append(ids, k)
+	}
+
 	if ids == nil {
 		if len(userIDs) == 0 && stored.Identity.Type == compute.ResourceIdentityTypeNone { // Validate that we have reset the resource type as none.
 			return true
 		}
 		return false
 	}
-	return reflect.DeepEqual(*ids, userIDs)
+
+	sort.Strings(ids)
+	sort.Strings(userIDs)
+	return reflect.DeepEqual(ids, userIDs)
 }
 
 func (c *TestCloudClient) ListMSI() (ret map[string]*[]string) {
@@ -324,22 +389,26 @@ func (c *TestCloudClient) UnSetError() {
 
 func NewTestVMClient() *TestVMClient {
 	nodeMap := make(map[string]*compute.VirtualMachine)
+	nodeIDs := make(map[string]map[string]bool)
 	vmClient := &VMClient{}
 
 	return &TestVMClient{
 		vmClient,
 		nodeMap,
+		nodeIDs,
 		nil,
 	}
 }
 
 func NewTestVMSSClient() *TestVMSSClient {
 	nodeMap := make(map[string]*compute.VirtualMachineScaleSet)
+	nodeIDs := make(map[string]map[string]bool)
 	vmssClient := &VMSSClient{}
 
 	return &TestVMSSClient{
 		vmssClient,
 		nodeMap,
+		nodeIDs,
 		nil,
 	}
 }
diff --git a/pkg/cloudprovider/identity_test.go b/pkg/cloudprovider/identity_test.go
index 28d0bc72..dadd869c 100644
--- a/pkg/cloudprovider/identity_test.go
+++ b/pkg/cloudprovider/identity_test.go
@@ -2,147 +2,8 @@ package cloudprovider
 
 import (
 	"testing"
-
-	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2018-04-01/compute"
 )
 
-func TestFilterIdentity(t *testing.T) {
-	idList := []string{}
-	idType := compute.ResourceIdentityTypeNone
-	if err := filterUserIdentity(&idType, &idList, "A"); err != nil {
-		t.Fatalf("expected nil error, got: %v", err)
-	}
-
-	idType = compute.ResourceIdentityTypeUserAssigned
-	if err := filterUserIdentity(&idType, &idList, "A"); err == nil || err != errNotAssigned {
-		t.Fatalf("expected error %q, got: %v", errNotAssigned, err)
-	}
-
-	idList = []string{"A"}
-	if err := filterUserIdentity(&idType, &idList, "A"); err != nil {
-		t.Fatal(err)
-	}
-	expect := []string{}
-	checkIDList(t, expect, idList)
-	if idType != compute.ResourceIdentityTypeNone {
-		t.Fatalf("expected id type to be %q, got: %s", compute.ResourceIdentityTypeNone, idType)
-	}
-
-	idList = []string{"A", "B"}
-	idType = compute.ResourceIdentityTypeUserAssigned
-	if err := filterUserIdentity(&idType, &idList, "A"); err != nil {
-		t.Fatal(err)
-	}
-	expect = []string{"B"}
-	checkIDList(t, expect, idList)
-	if idType != compute.ResourceIdentityTypeUserAssigned {
-		t.Fatalf("expected id type to be %q, got: %s", compute.ResourceIdentityTypeNone, idType)
-	}
-
-	idList = []string{"A", "B"}
-	idType = compute.ResourceIdentityTypeSystemAssignedUserAssigned
-	if err := filterUserIdentity(&idType, &idList, "A"); err != nil {
-		t.Fatal(err)
-	}
-	checkIDList(t, expect, idList)
-	if idType != compute.ResourceIdentityTypeSystemAssignedUserAssigned {
-		t.Fatalf("expected id type to be %q, got: %s", compute.ResourceIdentityTypeSystemAssignedUserAssigned, idType)
-	}
-
-	idList = []string{"A"}
-	idType = compute.ResourceIdentityTypeSystemAssignedUserAssigned
-	if err := filterUserIdentity(&idType, &idList, "A"); err != nil {
-		t.Fatal(err)
-	}
-	expect = []string{}
-	checkIDList(t, expect, idList)
-	if idType != compute.ResourceIdentityTypeSystemAssigned {
-		t.Fatalf("expected id type to be %q, got: %s", compute.ResourceIdentityTypeSystemAssigned, idType)
-	}
-}
-
-func TestAppendUserIdentity(t *testing.T) {
-	var (
-		idList []string
-		idType compute.ResourceIdentityType
-	)
-
-	append := appendUserIdentity(&idType, &idList, "A")
-	if !append {
-		t.Fatalf("Expecting the id to be not present. But present returned by Append.")
-	}
-	expect := []string{"A"}
-	checkIDList(t, expect, idList)
-	if idType != compute.ResourceIdentityTypeUserAssigned {
-		t.Fatalf("expected type %s, got: %s", compute.ResourceIdentityTypeUserAssigned, idType)
-	}
-
-	// Append the same value again, should not change anything
-	append = appendUserIdentity(&idType, &idList, "A")
-	if append {
-		t.Fatalf("Expecting the id to be not present. But present returned by Append.")
-	}
-	checkIDList(t, expect, idList)
-	if idType != compute.ResourceIdentityTypeUserAssigned {
-		t.Fatalf("expected type %s, got: %s", compute.ResourceIdentityTypeUserAssigned, idType)
-	}
-
-	append = appendUserIdentity(&idType, &idList, "B")
-	if !append {
-		t.Fatalf("Expecting the id to be not present. But present returned by Append.")
-	}
-	expect = []string{"A", "B"}
-	checkIDList(t, expect, idList)
-	if idType != compute.ResourceIdentityTypeUserAssigned {
-		t.Fatalf("expected type %s, got: %s", compute.ResourceIdentityTypeUserAssigned, idType)
-	}
-
-	idType = compute.ResourceIdentityTypeSystemAssigned
-	idList = []string{"A"}
-	expect = []string{"A", "B"}
-	append = appendUserIdentity(&idType, &idList, "B")
-	if !append {
-		t.Fatalf("Expecting the id to be not present. But present returned by Append.")
-	}
-	checkIDList(t, expect, idList)
-	if idType != compute.ResourceIdentityTypeSystemAssignedUserAssigned {
-		t.Fatalf("expected type %s, got: %s", compute.ResourceIdentityTypeSystemAssignedUserAssigned, idType)
-	}
-
-	idType = compute.ResourceIdentityTypeNone
-	idList = []string{}
-	expect = []string{"A"}
-	append = appendUserIdentity(&idType, &idList, "A")
-	if !append {
-		t.Fatalf("Expecting the id to be not present. But present returned by Append.")
-	}
-	checkIDList(t, expect, idList)
-	if idType != compute.ResourceIdentityTypeUserAssigned {
-		t.Fatalf("expected type %s, got: %s", compute.ResourceIdentityTypeUserAssigned, idType)
-	}
-
-	// test case sensitivity
-	idList = []string{}
-	expect = []string{"/subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/devcluster/providers/Microsoft.ManagedIdentity/userAssignedIdentities/keyvault-identity"}
-	append = appendUserIdentity(&idType, &idList, "/subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/devcluster/providers/Microsoft.ManagedIdentity/userAssignedIdentities/keyvault-identity")
-	if !append {
-		t.Fatalf("Expecting the id to be not present. But present returned by Append.")
-	}
-	checkIDList(t, expect, idList)
-	if idType != compute.ResourceIdentityTypeUserAssigned {
-		t.Fatalf("expected type %s, got: %s", compute.ResourceIdentityTypeUserAssigned, idType)
-	}
-	// append same id with non camel case resourcegroups
-	append = appendUserIdentity(&idType, &idList, "/subscriptions/00000000-0000-0000-0000-000000000000/resourcegroups/devcluster/providers/Microsoft.ManagedIdentity/userAssignedIdentities/keyvault-identity")
-	if append {
-		t.Fatalf("Expecting the id to be present. But not present returned by Append.")
-	}
-	checkIDList(t, expect, idList)
-	if idType != compute.ResourceIdentityTypeUserAssigned {
-		t.Fatalf("expected type %s, got: %s", compute.ResourceIdentityTypeUserAssigned, idType)
-	}
-}
-
 func checkIDList(t *testing.T, expect, actual []string) {
 	t.Helper()
 	if len(actual) != len(expect) {
diff --git a/pkg/cloudprovider/vmss_test.go b/pkg/cloudprovider/vmss_test.go
new file mode 100644
index 00000000..c697ede0
--- /dev/null
+++ b/pkg/cloudprovider/vmss_test.go
@@ -0,0 +1,24 @@
+package cloudprovider
+
+import (
+	"testing"
+
+	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2019-12-01/compute"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestSetUserIdentities(t *testing.T) {
+	testIdentityInfo := &vmssIdentityInfo{
+		info: &compute.VirtualMachineScaleSetIdentity{},
+	}
+
+	// adding id1
+	update := testIdentityInfo.SetUserIdentities(map[string]bool{"id1": true})
+	assert.True(t, update)
+	// adding id2
+	update = testIdentityInfo.SetUserIdentities(map[string]bool{"id2": true})
+	assert.True(t, update)
+	// add id3 and delete id1
+	update = testIdentityInfo.SetUserIdentities(map[string]bool{"id3": true, "id4": true, "id1": false})
+	assert.True(t, update)
+}
diff --git a/pkg/mic/mic_test.go b/pkg/mic/mic_test.go
index fc02361c..1f674f8e 100644
--- a/pkg/mic/mic_test.go
+++ b/pkg/mic/mic_test.go
@@ -4,6 +4,7 @@ import (
 	"errors"
 	"fmt"
 	"reflect"
+	"sort"
 	"sync"
 	"testing"
 	"time"
@@ -14,7 +15,7 @@ import (
 	"github.com/Azure/aad-pod-identity/pkg/crd"
 	"github.com/Azure/aad-pod-identity/pkg/metrics"
 
-	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2018-04-01/compute"
+	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2019-12-01/compute"
 
 	cp "github.com/Azure/aad-pod-identity/pkg/cloudprovider"
 	api "k8s.io/api/core/v1"
@@ -39,6 +40,7 @@ type TestVMClient struct {
 
 	mu       sync.Mutex
 	nodeMap  map[string]*compute.VirtualMachine
+	nodeIDs  map[string]map[string]bool
 	err      *error
 	identity *compute.VirtualMachineIdentity
 }
@@ -63,33 +65,58 @@ func (c *TestVMClient) Get(rgName string, nodeName string) (ret compute.VirtualM
 	if stored == nil {
 		vm := new(compute.VirtualMachine)
 		c.nodeMap[nodeName] = vm
+		c.nodeIDs[nodeName] = make(map[string]bool)
 		return *vm, nil
 	}
+
+	storedIDs := c.nodeIDs[nodeName]
+	newVMIdentity := make(map[string]*compute.VirtualMachineIdentityUserAssignedIdentitiesValue)
+	for id := range storedIDs {
+		newVMIdentity[id] = &compute.VirtualMachineIdentityUserAssignedIdentitiesValue{}
+	}
+	stored.Identity.UserAssignedIdentities = newVMIdentity
 	return *stored, nil
 }
 
-func (c *TestVMClient) CreateOrUpdate(rg string, nodeName string, vm compute.VirtualMachine) error {
+func (c *TestVMClient) UpdateIdentities(rg, nodeName string, vm compute.VirtualMachine) error {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
 	if c.err != nil {
-		c.nodeMap[nodeName].Identity = c.identity
 		return *c.err
 	}
+
+	if vm.Identity != nil && vm.Identity.UserAssignedIdentities != nil {
+		for k, v := range vm.Identity.UserAssignedIdentities {
+			if v == nil {
+				delete(c.nodeIDs[nodeName], k)
+			} else {
+				c.nodeIDs[nodeName][k] = true
+			}
+		}
+	}
+	if vm.Identity != nil && vm.Identity.UserAssignedIdentities == nil {
+		for k := range c.nodeIDs[nodeName] {
+			delete(c.nodeIDs[nodeName], k)
+		}
+	}
+
 	c.nodeMap[nodeName] = &vm
 	return nil
 }
 
 func (c *TestVMClient) ListMSI() (ret map[string]*[]string) {
-	ret = make(map[string]*[]string)
-
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
+	ret = make(map[string]*[]string)
+
 	for key, val := range c.nodeMap {
-		if val.Identity != nil {
-			ret[key] = val.Identity.IdentityIds
+		var ids []string
+		for k := range val.Identity.UserAssignedIdentities {
+			ids = append(ids, k)
 		}
+		ret[key] = &ids
 	}
 	return ret
 }
@@ -103,14 +130,20 @@ func (c *TestVMClient) CompareMSI(nodeName string, userIDs []string) bool {
 		return false
 	}
 
-	ids := stored.Identity.IdentityIds
+	var ids []string
+	for k := range c.nodeIDs[nodeName] {
+		ids = append(ids, k)
+	}
 	if ids == nil {
 		if len(userIDs) == 0 && stored.Identity.Type == compute.ResourceIdentityTypeNone { // Validate that we have reset the resource type as none.
 			return true
 		}
 		return false
 	}
-	return reflect.DeepEqual(*ids, userIDs)
+
+	sort.Strings(ids)
+	sort.Strings(userIDs)
+	return reflect.DeepEqual(ids, userIDs)
 }
 
 type TestVMSSClient struct {
@@ -118,6 +151,7 @@ type TestVMSSClient struct {
 
 	mu       sync.Mutex
 	nodeMap  map[string]*compute.VirtualMachineScaleSet
+	nodeIDs  map[string]map[string]bool
 	err      *error
 	identity *compute.VirtualMachineScaleSetIdentity
 }
@@ -138,31 +172,54 @@ func (c *TestVMSSClient) Get(rgName string, nodeName string) (ret compute.Virtua
 	if stored == nil {
 		vm := new(compute.VirtualMachineScaleSet)
 		c.nodeMap[nodeName] = vm
+		c.nodeIDs[nodeName] = make(map[string]bool)
 		return *vm, nil
 	}
+
+	storedIDs := c.nodeIDs[nodeName]
+	newVMSSIdentity := make(map[string]*compute.VirtualMachineScaleSetIdentityUserAssignedIdentitiesValue)
+	for id := range storedIDs {
+		newVMSSIdentity[id] = &compute.VirtualMachineScaleSetIdentityUserAssignedIdentitiesValue{}
+	}
+	stored.Identity.UserAssignedIdentities = newVMSSIdentity
 	return *stored, nil
 }
 
-func (c *TestVMSSClient) CreateOrUpdate(rg string, nodeName string, vm compute.VirtualMachineScaleSet) error {
+func (c *TestVMSSClient) UpdateIdentities(rg, nodeName string, vmss compute.VirtualMachineScaleSet) error {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
 	if c.err != nil {
-		c.nodeMap[nodeName].Identity = c.identity
 		return *c.err
 	}
-	c.nodeMap[nodeName] = &vm
+	if vmss.Identity != nil && vmss.Identity.UserAssignedIdentities != nil {
+		for k, v := range vmss.Identity.UserAssignedIdentities {
+			if v == nil {
+				delete(c.nodeIDs[nodeName], k)
+			} else {
+				c.nodeIDs[nodeName][k] = true
+			}
+		}
+	}
+	if vmss.Identity != nil && vmss.Identity.UserAssignedIdentities == nil {
+		for k := range c.nodeIDs[nodeName] {
+			delete(c.nodeIDs[nodeName], k)
+		}
+	}
+
+	c.nodeMap[nodeName] = &vmss
 	return nil
 }
 
 func (c *TestVMSSClient) ListMSI() (ret map[string]*[]string) {
 	ret = make(map[string]*[]string)
 
-	c.mu.Lock()
-	defer c.mu.Unlock()
-
 	for key, val := range c.nodeMap {
-		ret[key] = val.Identity.IdentityIds
+		var ids []string
+		for k := range val.Identity.UserAssignedIdentities {
+			ids = append(ids, k)
+		}
+		ret[key] = &ids
 	}
 	return ret
 }
@@ -176,14 +233,21 @@ func (c *TestVMSSClient) CompareMSI(nodeName string, userIDs []string) bool {
 		return false
 	}
 
-	ids := stored.Identity.IdentityIds
+	var ids []string
+	for k := range c.nodeIDs[nodeName] {
+		ids = append(ids, k)
+	}
+
 	if ids == nil {
 		if len(userIDs) == 0 && stored.Identity.Type == compute.ResourceIdentityTypeNone { // Validate that we have reset the resource type as none.
 			return true
 		}
 		return false
 	}
-	return reflect.DeepEqual(*ids, userIDs)
+
+	sort.Strings(ids)
+	sort.Strings(userIDs)
+	return reflect.DeepEqual(ids, userIDs)
 }
 
 func (c *TestCloudClient) ListMSI() (ret map[string]*[]string) {
@@ -221,24 +285,32 @@ func (c *TestCloudClient) UnSetError() {
 
 func NewTestVMClient() *TestVMClient {
 	nodeMap := make(map[string]*compute.VirtualMachine)
+	nodeIDs := make(map[string]map[string]bool)
 	vmClient := &cp.VMClient{}
-	identity := &compute.VirtualMachineIdentity{IdentityIds: &[]string{}}
+	identity := &compute.VirtualMachineIdentity{
+		UserAssignedIdentities: make(map[string]*compute.VirtualMachineIdentityUserAssignedIdentitiesValue),
+	}
 
 	return &TestVMClient{
 		VMClient: vmClient,
 		nodeMap:  nodeMap,
+		nodeIDs:  nodeIDs,
 		identity: identity,
 	}
 }
 
 func NewTestVMSSClient() *TestVMSSClient {
 	nodeMap := make(map[string]*compute.VirtualMachineScaleSet)
+	nodeIDs := make(map[string]map[string]bool)
 	vmssClient := &cp.VMSSClient{}
-	identity := &compute.VirtualMachineScaleSetIdentity{IdentityIds: &[]string{}}
+	identity := &compute.VirtualMachineScaleSetIdentity{
+		UserAssignedIdentities: make(map[string]*compute.VirtualMachineScaleSetIdentityUserAssignedIdentitiesValue),
+	}
 
 	return &TestVMSSClient{
 		VMSSClient: vmssClient,
 		nodeMap:    nodeMap,
+		nodeIDs:    nodeIDs,
 		identity:   identity,
 	}
 }
@@ -1073,7 +1145,11 @@ func TestMICStateFlow(t *testing.T) {
 	podClient.DeletePod("test-pod1", "default")
 	// SetError sets error in crd client only for remove assigned identity
 	cloudClient.SetError(errors.New("error removing identity from node"))
-	cloudClient.testVMClient.identity = &compute.VirtualMachineIdentity{IdentityIds: &[]string{"test-user-msi-resourceid"}}
+	cloudClient.testVMClient.identity = &compute.VirtualMachineIdentity{
+		UserAssignedIdentities: map[string]*compute.VirtualMachineIdentityUserAssignedIdentitiesValue{
+			"test-user-msi-resourceid": &compute.VirtualMachineIdentityUserAssignedIdentitiesValue{},
+		},
+	}
 
 	eventCh <- internalaadpodid.PodDeleted
 	if !evtRecorder.WaitForEvents(1) {
@@ -1254,7 +1330,11 @@ func TestSyncRetryLoop(t *testing.T) {
 	// delete the pod, simulate failure in cloud calls on trying to un-assign identity from node
 	podClient.DeletePod("test-pod1", "default")
 	cloudClient.SetError(errors.New("error removing identity from node"))
-	cloudClient.testVMClient.identity = &compute.VirtualMachineIdentity{IdentityIds: &[]string{"test-user-msi-resourceid"}}
+	cloudClient.testVMClient.identity = &compute.VirtualMachineIdentity{
+		UserAssignedIdentities: map[string]*compute.VirtualMachineIdentityUserAssignedIdentitiesValue{
+			"test-user-msi-resourceid": &compute.VirtualMachineIdentityUserAssignedIdentitiesValue{},
+		},
+	}
 
 	eventCh <- internalaadpodid.PodDeleted
 	if !evtRecorder.WaitForEvents(1) {
diff --git a/test/e2e/identityvalidator/identityvalidator.go b/test/e2e/identityvalidator/identityvalidator.go
index 7adcf61c..8e7fcfa2 100644
--- a/test/e2e/identityvalidator/identityvalidator.go
+++ b/test/e2e/identityvalidator/identityvalidator.go
@@ -8,7 +8,7 @@ import (
 	"github.com/pkg/errors"
 	"k8s.io/klog"
 
-	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2018-04-01/compute"
+	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2019-12-01/compute"
 	"github.com/Azure/azure-sdk-for-go/services/keyvault/2016-10-01/keyvault"
 	"github.com/Azure/azure-sdk-for-go/services/keyvault/auth"
 	"github.com/Azure/go-autorest/autorest"