Repository: Azure/aad-pod-identity
Pull request #559 (2020-04-10 06:28:41)

fix updating assigned identity when azure identity updated
<!-- 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? -->
- Adds code path to update azure assigned identity when azure identity is updated in place
- Ensures old identity is deleted and new updated identity is assigned
- Adds e2e test to test update path

**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/510, fixes https://github.com/Azure/aad-pod-identity/issues/504

**Notes for Reviewers**:



Patch:
diff --git a/pkg/cloudprovider/cloudprovider.go b/pkg/cloudprovider/cloudprovider.go
index fce716a9..df225c05 100644
--- a/pkg/cloudprovider/cloudprovider.go
+++ b/pkg/cloudprovider/cloudprovider.go
@@ -191,7 +191,7 @@ func (c *Client) UpdateUserMSI(addUserAssignedMSIIDs, removeUserAssignedMSIIDs [
 	requiresUpdate := info.SetUserIdentities(ids)
 
 	if requiresUpdate {
-		klog.Infof("Updating user assigned MSIs on %s", name)
+		klog.Infof("Updating user assigned MSIs on %s, assign [%d], unassign [%d]", name, len(addUserAssignedMSIIDs), len(removeUserAssignedMSIIDs))
 		timeStarted := time.Now()
 		if err := updateFunc(); err != nil {
 			return err
diff --git a/pkg/cloudprovider/vm.go b/pkg/cloudprovider/vm.go
index 94a608c2..bc6b761b 100644
--- a/pkg/cloudprovider/vm.go
+++ b/pkg/cloudprovider/vm.go
@@ -144,7 +144,7 @@ func (i *vmIdentityInfo) SetUserIdentities(ids map[string]bool) bool {
 
 	nodeList := make(map[string]bool)
 	// add all current existing ids
-	for id, _ := range i.info.UserAssignedIdentities {
+	for id := range i.info.UserAssignedIdentities {
 		nodeList[id] = true
 	}
 
diff --git a/pkg/cloudprovider/vmss.go b/pkg/cloudprovider/vmss.go
index 86ebf3ea..973e0cba 100644
--- a/pkg/cloudprovider/vmss.go
+++ b/pkg/cloudprovider/vmss.go
@@ -146,7 +146,7 @@ func (i *vmssIdentityInfo) SetUserIdentities(ids map[string]bool) bool {
 
 	nodeList := make(map[string]bool)
 	// add all current existing ids
-	for id, _ := range i.info.UserAssignedIdentities {
+	for id := range i.info.UserAssignedIdentities {
 		nodeList[id] = true
 	}
 
diff --git a/pkg/crd/crd.go b/pkg/crd/crd.go
index 4882cead..a4c79497 100644
--- a/pkg/crd/crd.go
+++ b/pkg/crd/crd.go
@@ -41,6 +41,7 @@ type ClientInt interface {
 	SyncCacheAll(exit <-chan struct{}, initial bool)
 	RemoveAssignedIdentity(assignedIdentity *aadpodid.AzureAssignedIdentity) error
 	CreateAssignedIdentity(assignedIdentity *aadpodid.AzureAssignedIdentity) error
+	UpdateAssignedIdentity(assignedIdentity *aadpodid.AzureAssignedIdentity) error
 	UpdateAzureAssignedIdentityStatus(assignedIdentity *aadpodid.AzureAssignedIdentity, status string) error
 	UpgradeAll() error
 	ListBindings() (res *[]aadpodid.AzureIdentityBinding, err error)
@@ -519,6 +520,32 @@ func (c *Client) CreateAssignedIdentity(assignedIdentity *aadpodid.AzureAssigned
 	return nil
 }
 
+func (c *Client) UpdateAssignedIdentity(assignedIdentity *aadpodid.AzureAssignedIdentity) (err error) {
+	klog.Infof("Got assigned id %s/%s to update", assignedIdentity.Namespace, assignedIdentity.Name)
+	begin := time.Now()
+
+	defer func() {
+		if err != nil {
+			c.reporter.ReportKubernetesAPIOperationError(metrics.AssignedIdentityUpdateOperationName)
+			return
+		}
+		c.reporter.Report(
+			metrics.AssignedIdentityUpdateCountM.M(1),
+			metrics.AssignedIdentityUpdateDurationM.M(metrics.SinceInSeconds(begin)))
+
+	}()
+
+	v1AssignedID := aadpodv1.ConvertInternalAssignedIdentityToV1AssignedIdentity(*assignedIdentity)
+	err = c.rest.Put().Namespace(assignedIdentity.Namespace).Resource(aadpodid.AzureAssignedIDResource).Name(assignedIdentity.Name).Body(&v1AssignedID).Do().Error()
+	if err != nil {
+		return fmt.Errorf("failed to update assigned identity: %v", err)
+	}
+
+	klog.V(5).Infof("Time take to update %s/%s: %v", assignedIdentity.Namespace, assignedIdentity.Name, time.Since(begin))
+	stats.Update(stats.AssignedIDAdd, time.Since(begin))
+	return nil
+}
+
 // ListBindings returns a list of azureidentitybindings
 func (c *Client) ListBindings() (res *[]aadpodid.AzureIdentityBinding, err error) {
 	begin := time.Now()
diff --git a/pkg/metrics/metrics.go b/pkg/metrics/metrics.go
index 712c334a..a5aad09c 100644
--- a/pkg/metrics/metrics.go
+++ b/pkg/metrics/metrics.go
@@ -18,6 +18,8 @@ const (
 	assignedIdentityAdditionCountName      = "assigned_identity_addition_count"
 	assignedIdentityDeletionDurationName   = "assigned_identity_deletion_duration_seconds"
 	assignedIdentityDeletionCountName      = "assigned_identity_deletion_count"
+	assignedIdentityUpdateDurationName     = "assigned_identity_update_duration_seconds"
+	assignedIdentityUpdateCountName        = "assigned_identity_update_count"
 	nmiOperationsDurationName              = "nmi_operations_duration_seconds"
 	micCycleDurationName                   = "mic_cycle_duration_seconds"
 	micCycleCountName                      = "mic_cycle_count"
@@ -46,6 +48,8 @@ const (
 	AssignedIdentityDeletionOperationName = "assigned_identity_deletion"
 	// AssignedIdentityAdditionOperationName ...
 	AssignedIdentityAdditionOperationName = "assigned_identity_addition"
+	// AssignedIdentityUpdateOperationName ...
+	AssignedIdentityUpdateOperationName = "assigned_identity_update"
 	// UpdateAzureAssignedIdentityStatusOperationName ...
 	UpdateAzureAssignedIdentityStatusOperationName = "update_azure_assigned_identity_status"
 	// GetPodListOperationName
@@ -132,6 +136,18 @@ var (
 		imdsOperationsDurationName,
 		"Duration in seconds of imds token operations",
 		stats.UnitMilliseconds)
+
+	// AssignedIdentityUpdateDurationM is a measure that tracks the duration in seconds of assigned_identity_update operations.
+	AssignedIdentityUpdateDurationM = stats.Float64(
+		assignedIdentityUpdateDurationName,
+		"Duration in seconds of assigned identity update operations",
+		stats.UnitMilliseconds)
+
+	// AssignedIdentityUpdateCountM is a measure that tracks the cumulative number of assigned identity update operations.
+	AssignedIdentityUpdateCountM = stats.Int64(
+		assignedIdentityUpdateCountName,
+		"Total number of assigned identity update operations",
+		stats.UnitDimensionless)
 )
 
 var (
@@ -222,6 +238,16 @@ func registerViews() error {
 			Aggregation: view.Distribution(0.01, 0.02, 0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1, 2, 3, 4, 5, 10),
 			TagKeys:     []tag.Key{operationTypeKey},
 		},
+		&view.View{
+			Description: AssignedIdentityUpdateDurationM.Description(),
+			Measure:     AssignedIdentityUpdateDurationM,
+			Aggregation: view.Distribution(0.01, 0.02, 0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1, 2, 3, 4, 5, 10),
+		},
+		&view.View{
+			Description: AssignedIdentityUpdateCountM.Description(),
+			Measure:     AssignedIdentityUpdateCountM,
+			Aggregation: view.Count(),
+		},
 	}
 	err := view.Register(views...)
 	return err
diff --git a/pkg/mic/mic.go b/pkg/mic/mic.go
index d4e02caa..3a7338ab 100644
--- a/pkg/mic/mic.go
+++ b/pkg/mic/mic.go
@@ -115,6 +115,7 @@ type trackUserAssignedMSIIds struct {
 	removeUserAssignedMSIIDs []string
 	assignedIDsToCreate      []aadpodid.AzureAssignedIdentity
 	assignedIDsToDelete      []aadpodid.AzureAssignedIdentity
+	assignedIDsToUpdate      []aadpodid.AzureAssignedIdentity
 	isvmss                   bool
 }
 
@@ -428,24 +429,25 @@ func (c *Client) Sync(exit <-chan struct{}) {
 			klog.Error(err)
 			continue
 		}
-		klog.V(5).Infof("del: %v, add: %v", deleteList, addList)
+		beforeUpdateList, afterUpdateList := c.getAzureAssignedIdentitiesToUpdate(addList, deleteList)
+		klog.V(5).Infof("del: %v, add: %v, update: %v", deleteList, addList, afterUpdateList)
 
 		// the node map is used to track assigned ids to create/delete, identities to assign/remove
 		// for each node or vmss
 		nodeMap := make(map[string]trackUserAssignedMSIIds)
 
-		// seperate the add and delete list per node
-		c.convertAssignedIDListToMap(addList, deleteList, nodeMap)
+		// separate the add, delete and update list per node
+		c.convertAssignedIDListToMap(addList, deleteList, afterUpdateList, nodeMap)
 
 		// process the delete and add list
 		// determine the list of identities that need to updated, create a node to identity list mapping for add and delete
-		if len(deleteList) > 0 {
+		if len(deleteList) > 0 || len(beforeUpdateList) > 0 {
 			workDone = true
-			c.getListOfIdsToDelete(deleteList, newAssignedIDs, nodeMap, nodeRefs)
+			c.getListOfIdsToDelete(deleteList, beforeUpdateList, afterUpdateList, newAssignedIDs, nodeMap, nodeRefs)
 		}
-		if len(addList) > 0 {
+		if len(addList) > 0 || len(afterUpdateList) > 0 {
 			workDone = true
-			c.getListOfIdsToAssign(addList, nodeMap)
+			c.getListOfIdsToAssign(addList, afterUpdateList, nodeMap)
 		}
 
 		var wg sync.WaitGroup
@@ -490,7 +492,7 @@ func (c *Client) Sync(exit <-chan struct{}) {
 	}
 }
 
-func (c *Client) convertAssignedIDListToMap(addList, deleteList map[string]aadpodid.AzureAssignedIdentity, nodeMap map[string]trackUserAssignedMSIIds) {
+func (c *Client) convertAssignedIDListToMap(addList, deleteList, updateList map[string]aadpodid.AzureAssignedIdentity, nodeMap map[string]trackUserAssignedMSIIds) {
 	if addList != nil {
 		for _, createID := range addList {
 			if trackList, ok := nodeMap[createID.Spec.NodeName]; ok {
@@ -512,6 +514,17 @@ func (c *Client) convertAssignedIDListToMap(addList, deleteList map[string]aadpo
 			nodeMap[delID.Spec.NodeName] = trackUserAssignedMSIIds{assignedIDsToDelete: []aadpodid.AzureAssignedIdentity{delID}}
 		}
 	}
+
+	if updateList != nil {
+		for _, updateID := range updateList {
+			if trackList, ok := nodeMap[updateID.Spec.NodeName]; ok {
+				trackList.assignedIDsToUpdate = append(trackList.assignedIDsToUpdate, updateID)
+				nodeMap[updateID.Spec.NodeName] = trackList
+				continue
+			}
+			nodeMap[updateID.Spec.NodeName] = trackUserAssignedMSIIds{assignedIDsToUpdate: []aadpodid.AzureAssignedIdentity{updateID}}
+		}
+	}
 }
 
 func (c *Client) createDesiredAssignedIdentityList(
@@ -581,8 +594,7 @@ func (c *Client) createDesiredAssignedIdentityList(
 
 // getListOfIdsToDelete will go over the delete list to determine if the id is required to be deleted
 // only user assigned identity not in use are added to the remove list for the node
-func (c *Client) getListOfIdsToDelete(deleteList map[string]aadpodid.AzureAssignedIdentity,
-	newAssignedIDs map[string]aadpodid.AzureAssignedIdentity,
+func (c *Client) getListOfIdsToDelete(deleteList, beforeUpdateList, afterUpdateList, newAssignedIDs map[string]aadpodid.AzureAssignedIdentity,
 	nodeMap map[string]trackUserAssignedMSIIds,
 	nodeRefs map[string]bool) {
 	vmssGroups, err := getVMSSGroups(c.NodeClient, nodeRefs)
@@ -590,43 +602,75 @@ func (c *Client) getListOfIdsToDelete(deleteList map[string]aadpodid.AzureAssign
 		klog.Error(err)
 		return
 	}
+
+	consolidatedMapToCheck := make(map[string]aadpodid.AzureAssignedIdentity)
+	for name, id := range newAssignedIDs {
+		consolidatedMapToCheck[name] = id
+	}
+	for name, id := range afterUpdateList {
+		consolidatedMapToCheck[name] = id
+	}
+
 	for _, delID := range deleteList {
-		klog.V(5).Infof("Deletion of id: %s", delID.Name)
-		inUse, err := c.checkIfInUse(delID, newAssignedIDs, vmssGroups)
+		err := c.shouldRemoveID(delID, consolidatedMapToCheck, nodeMap, vmssGroups)
 		if err != nil {
 			klog.Error(err)
-			continue
 		}
-
-		id := delID.Spec.AzureIdentityRef
-		isUserAssignedMSI := c.checkIfUserAssignedMSI(*id)
-		isImmutableIdentity := c.checkIfIdentityImmutable(id.Spec.ClientID)
-
-		// this case includes Assigned state and empty state to ensure backward compatability
-		if delID.Status.Status == aadpodid.AssignedIDAssigned || delID.Status.Status == "" {
-			// only user assigned identities that are not in use and are not defined as
-			// immutable will be removed from underlying node/vmss
-			if !inUse && isUserAssignedMSI && !isImmutableIdentity {
-				c.appendToRemoveListForNode(id.Spec.ResourceID, delID.Spec.NodeName, nodeMap)
-			}
+	}
+	// this loop checks the azure identity before it was updated and cleans up
+	// the old identity
+	for _, oldUpdateID := range beforeUpdateList {
+		err := c.shouldRemoveID(oldUpdateID, consolidatedMapToCheck, nodeMap, vmssGroups)
+		if err != nil {
+			klog.Error(err)
 		}
-		klog.V(5).Infof("Binding removed: %+v", delID.Spec.AzureBindingRef)
 	}
 }
 
 // getListOfIdsToAssign will add the id to the append list for node if it's user assigned identity
-func (c *Client) getListOfIdsToAssign(addList map[string]aadpodid.AzureAssignedIdentity, nodeMap map[string]trackUserAssignedMSIIds) {
+func (c *Client) getListOfIdsToAssign(addList, updateList map[string]aadpodid.AzureAssignedIdentity, nodeMap map[string]trackUserAssignedMSIIds) {
 	for _, createID := range addList {
-		id := createID.Spec.AzureIdentityRef
-		isUserAssignedMSI := c.checkIfUserAssignedMSI(*id)
+		c.shouldAssignID(createID, nodeMap)
+	}
+	for _, updateID := range updateList {
+		c.shouldAssignID(updateID, nodeMap)
+	}
+}
 
-		if createID.Status.Status == "" || createID.Status.Status == aadpodid.AssignedIDCreated {
-			if isUserAssignedMSI {
-				c.appendToAddListForNode(id.Spec.ResourceID, createID.Spec.NodeName, nodeMap)
-			}
+func (c *Client) shouldAssignID(assignedID aadpodid.AzureAssignedIdentity, nodeMap map[string]trackUserAssignedMSIIds) {
+	id := assignedID.Spec.AzureIdentityRef
+	isUserAssignedMSI := c.checkIfUserAssignedMSI(*id)
+
+	if assignedID.Status.Status == "" || assignedID.Status.Status == aadpodid.AssignedIDCreated {
+		if isUserAssignedMSI {
+			c.appendToAddListForNode(id.Spec.ResourceID, assignedID.Spec.NodeName, nodeMap)
 		}
-		klog.V(5).Infof("Binding applied: %+v", createID.Spec.AzureBindingRef)
 	}
+	klog.V(5).Infof("Binding applied: %+v", assignedID.Spec.AzureBindingRef)
+}
+
+func (c *Client) shouldRemoveID(assignedID aadpodid.AzureAssignedIdentity,
+	newAssignedIDs map[string]aadpodid.AzureAssignedIdentity,
+	nodeMap map[string]trackUserAssignedMSIIds, vmssGroups *vmssGroupList) error {
+	klog.V(5).Infof("Deletion of id: %s", assignedID.Name)
+	inUse, err := c.checkIfInUse(assignedID, newAssignedIDs, vmssGroups)
+	if err != nil {
+		return err
+	}
+
+	id := assignedID.Spec.AzureIdentityRef
+	isUserAssignedMSI := c.checkIfUserAssignedMSI(*id)
+	isImmutableIdentity := c.checkIfIdentityImmutable(id.Spec.ClientID)
+	// this case includes Assigned state and empty state to ensure backward compatability
+	if assignedID.Status.Status == aadpodid.AssignedIDAssigned || assignedID.Status.Status == "" {
+		// only user assigned identities that are not in use and are not defined as
+		// immutable will be removed from underlying node/vmss
+		if !inUse && isUserAssignedMSI && !isImmutableIdentity {
+			c.appendToRemoveListForNode(id.Spec.ResourceID, assignedID.Spec.NodeName, nodeMap)
+		}
+	}
+	klog.V(5).Infof("Binding removed: %+v", assignedID.Spec.AzureBindingRef)
+	return nil
 }
 
 func (c *Client) matchAssignedID(x aadpodid.AzureAssignedIdentity, y aadpodid.AzureAssignedIdentity) (ret bool) {
@@ -673,7 +717,7 @@ func (c *Client) getAzureAssignedIDsToCreate(old, new map[string]aadpodid.AzureA
 			// if the old assigned id is in created state, then the identity assignment to the node
 			// is not done. Adding to the list will ensure we retry identity assignment to node for
 			// this assigned identity.
-			klog.V(5).Infof("ok: %v, Create added: %s", idMatch, assignedIDName)
+			klog.V(5).Infof("ok: %v, Create added: %s as assignedID in CREATED state", idMatch, assignedIDName)
 			create[assignedIDName] = oldAssignedID
 		}
 		if !idMatch {
@@ -718,6 +762,31 @@ func (c *Client) getAzureAssignedIDsToDelete(old, new map[string]aadpodid.AzureA
 	return delete, nil
 }
 
+// getAzureAssignedIdentitiesToUpdate returns a list of assignedIDs that need to be updated
+// because of change in azureIdentity or azurerIdentityBinding
+// returns 2 maps, first the assigned IDs currently on cluster, second the assignedID value to update with
+func (c *Client) getAzureAssignedIdentitiesToUpdate(add, del map[string]aadpodid.AzureAssignedIdentity) (map[string]aadpodid.AzureAssignedIdentity, map[string]aadpodid.AzureAssignedIdentity) {
+	beforeUpdate := make(map[string]aadpodid.AzureAssignedIdentity)
+	afterUpdate := make(map[string]aadpodid.AzureAssignedIdentity)
+	// no updates required as assigned identities will not be in both lists
+	if len(add) == 0 || len(del) == 0 {
+		return beforeUpdate, afterUpdate
+	}
+	for assignedIDName, addAssignedID := range add {
+		if delAssignedID, exists := del[assignedIDName]; exists {
+			// assigned identity exists in add and del list
+			// update the assigned identity to the latest
+			addAssignedID.ObjectMeta = delAssignedID.ObjectMeta
+			beforeUpdate[assignedIDName] = delAssignedID
+			afterUpdate[assignedIDName] = addAssignedID
+			// since this is part of update, remove the assignedID from the add and del list
+			delete(add, assignedIDName)
+			delete(del, assignedIDName)
+		}
+	}
+	return beforeUpdate, afterUpdate
+}
+
 func (c *Client) makeAssignedIDs(azID aadpodid.AzureIdentity, azBinding aadpodid.AzureIdentityBinding, podName, podNameSpace, nodeName string) (res *aadpodid.AzureAssignedIdentity, err error) {
 	binding := azBinding
 	id := azID
@@ -773,6 +842,10 @@ func (c *Client) removeAssignedIdentity(assignedID *aadpodid.AzureAssignedIdenti
 	return nil
 }
 
+func (c *Client) updateAssignedIdentity(assignedID *aadpodid.AzureAssignedIdentity) error {
+	return c.CRDClient.UpdateAssignedIdentity(assignedID)
+}
+
 func (c *Client) appendToRemoveListForNode(resourceID, nodeName string, nodeMap map[string]trackUserAssignedMSIIds) {
 	if trackList, ok := nodeMap[nodeName]; ok {
 		trackList.removeUserAssignedMSIIDs = append(trackList.removeUserAssignedMSIIDs, resourceID)
@@ -900,22 +973,23 @@ func (c *Client) updateNodeAndDeps(newAssignedIDs map[string]aadpodid.AzureAssig
 func (c *Client) updateUserMSI(newAssignedIDs map[string]aadpodid.AzureAssignedIdentity, nodeOrVMSSName string, nodeTrackList trackUserAssignedMSIIds, nodeRefs map[string]bool, wg *sync.WaitGroup) {
 	defer wg.Done()
 	beginAdding := time.Now()
-	klog.Infof("Processing node %s, add [%d], del [%d]", nodeOrVMSSName, len(nodeTrackList.assignedIDsToCreate), len(nodeTrackList.assignedIDsToDelete))
+	klog.Infof("Processing node %s, add [%d], del [%d], update [%d]", nodeOrVMSSName,
+		len(nodeTrackList.assignedIDsToCreate), len(nodeTrackList.assignedIDsToDelete), len(nodeTrackList.assignedIDsToUpdate))
 
 	ctx := context.TODO()
 	// We have to ensure that we don't overwhelm the API server with too many
 	// requests in flight. We use a token based approach implemented using semaphore to
 	// ensure that only given createDeleteBatch requests are in flight at any point in time.
 	// Note that at this point in the code path, we are doing this in parallel per node/VMSS already.
-	semCreate := semaphore.NewWeighted(c.createDeleteBatch)
+	semCreateOrUpdate := semaphore.NewWeighted(c.createDeleteBatch)
 
 	for _, createID := range nodeTrackList.assignedIDsToCreate {
-		if err := semCreate.Acquire(ctx, 1); err != nil {
+		if err := semCreateOrUpdate.Acquire(ctx, 1); err != nil {
 			klog.Errorf("Failed to acquire semaphore in the create loop: %v", err)
 			return
 		}
 		go func(assignedID aadpodid.AzureAssignedIdentity) {
-			defer semCreate.Release(1)
+			defer semCreateOrUpdate.Release(1)
 			if assignedID.Status.Status == "" {
 				binding := assignedID.Spec.AzureBindingRef
 
@@ -933,14 +1007,40 @@ func (c *Client) updateUserMSI(newAssignedIDs map[string]aadpodid.AzureAssignedI
 		}(createID)
 	}
 
+	for _, updateID := range nodeTrackList.assignedIDsToUpdate {
+		if err := semCreateOrUpdate.Acquire(ctx, 1); err != nil {
+			klog.Errorf("Failed to acquire semaphore in the update loop: %v", err)
+			return
+		}
+		go func(assignedID aadpodid.AzureAssignedIdentity) {
+			defer semCreateOrUpdate.Release(1)
+			if assignedID.Status.Status == "" {
+				binding := assignedID.Spec.AzureBindingRef
+
+				// this is the state when the azure assigned identity is yet to be created
+				klog.V(5).Infof("Initiating assigned id creation for pod - %s, binding - %s", assignedID.Spec.Pod, binding.Name)
+
+				assignedID.Status.Status = aadpodid.AssignedIDCreated
+				err := c.updateAssignedIdentity(&assignedID)
+				if err != nil {
+					c.EventRecorder.Event(binding, corev1.EventTypeWarning, "binding apply error",
+						fmt.Sprintf("Updating assigned identity for pod %s/%s resulted in error %v", assignedID.Namespace, assignedID.Name, err))
+					klog.Error(err)
+				}
+			}
+		}(updateID)
+	}
+
 	// Ensure that all creates are complete
-	if err := semCreate.Acquire(ctx, c.createDeleteBatch); err != nil {
+	if err := semCreateOrUpdate.Acquire(ctx, c.createDeleteBatch); err != nil {
 		klog.Errorf("Failed to acquire semaphore at the end of creates: %v", err)
 		return
 	}
 	// generate unique list so we don't make multiple calls to assign/remove same id
 	addUserAssignedMSIIDs := c.getUniqueIDs(nodeTrackList.addUserAssignedMSIIDs)
 	removeUserAssignedMSIIDs := c.getUniqueIDs(nodeTrackList.removeUserAssignedMSIIDs)
+	createOrUpdateList := append([]aadpodid.AzureAssignedIdentity{}, nodeTrackList.assignedIDsToCreate...)
+	createOrUpdateList = append(createOrUpdateList, nodeTrackList.assignedIDsToUpdate...)
 
 	err := c.CloudClient.UpdateUserMSI(addUserAssignedMSIIDs, removeUserAssignedMSIIDs, nodeOrVMSSName, nodeTrackList.isvmss)
 	if err != nil {
@@ -951,7 +1051,7 @@ func (c *Client) updateUserMSI(newAssignedIDs map[string]aadpodid.AzureAssignedI
 			return
 		}
 
-		for _, createID := range nodeTrackList.assignedIDsToCreate {
+		for _, createID := range createOrUpdateList {
 			id := createID.Spec.AzureIdentityRef
 			binding := createID.Spec.AzureBindingRef
 
@@ -1021,7 +1121,7 @@ func (c *Client) updateUserMSI(newAssignedIDs map[string]aadpodid.AzureAssignedI
 
 	semUpdate := semaphore.NewWeighted(c.createDeleteBatch)
 
-	for _, createID := range nodeTrackList.assignedIDsToCreate {
+	for _, createID := range createOrUpdateList {
 		if err := semUpdate.Acquire(ctx, 1); err != nil {
 			klog.Errorf("Failed to acquire semaphore in the update loop: %v", err)
 			return
@@ -1088,6 +1188,7 @@ func (c *Client) updateUserMSI(newAssignedIDs map[string]aadpodid.AzureAssignedI
 	}
 
 	stats.UpdateCount(stats.TotalAssignedIDsCreated, len(nodeTrackList.assignedIDsToCreate))
+	stats.UpdateCount(stats.TotalAssignedIDsUpdated, len(nodeTrackList.assignedIDsToUpdate))
 	stats.UpdateCount(stats.TotalAssignedIDsDeleted, len(nodeTrackList.assignedIDsToDelete))
 	stats.Put(stats.TotalCreateOrUpdate, time.Since(beginAdding))
 }
@@ -1154,16 +1255,15 @@ func (c *Client) consolidateVMSSNodes(nodeMap map[string]trackUserAssignedMSIIds
 		}
 
 		vmssTrackList := trackUserAssignedMSIIds{}
-
 		for _, vmssNode := range vmssNodes {
 			vmssTrackList.addUserAssignedMSIIDs = append(vmssTrackList.addUserAssignedMSIIDs, nodeMap[vmssNode].addUserAssignedMSIIDs...)
 			vmssTrackList.removeUserAssignedMSIIDs = append(vmssTrackList.removeUserAssignedMSIIDs, nodeMap[vmssNode].removeUserAssignedMSIIDs...)
 			vmssTrackList.assignedIDsToCreate = append(vmssTrackList.assignedIDsToCreate, nodeMap[vmssNode].assignedIDsToCreate...)
 			vmssTrackList.assignedIDsToDelete = append(vmssTrackList.assignedIDsToDelete, nodeMap[vmssNode].assignedIDsToDelete...)
+			vmssTrackList.assignedIDsToUpdate = append(vmssTrackList.assignedIDsToUpdate, nodeMap[vmssNode].assignedIDsToUpdate...)
 			vmssTrackList.isvmss = true
 
 			delete(nodeMap, vmssNode)
-
 			nodeMap[getVMSSName(vmssName)] = vmssTrackList
 		}
 	}
diff --git a/pkg/stats/stats.go b/pkg/stats/stats.go
index 19dc213b..19989831 100644
--- a/pkg/stats/stats.go
+++ b/pkg/stats/stats.go
@@ -29,6 +29,7 @@ const (
 	TotalPutCalls           StatsType = "Number of cloud provider PUT"
 	TotalGetCalls           StatsType = "Number of cloud provider GET"
 	TotalAssignedIDsCreated StatsType = "Number of assigned ids created in this sync cycle"
+	TotalAssignedIDsUpdated StatsType = "Number of assigned ids updated in this sync cycle"
 	TotalAssignedIDsDeleted StatsType = "Number of assigned ids deleted in this sync cycle"
 	K8sGet                  StatsType = "K8s get"
 	K8sPut                  StatsType = "K8s put"
@@ -114,6 +115,7 @@ func PrintSync() {
 		PrintCount(TotalGetCalls)
 
 		PrintCount(TotalAssignedIDsCreated)
+		PrintCount(TotalAssignedIDsUpdated)
 		PrintCount(TotalAssignedIDsDeleted)
 
 		Print(FindAssignedIDCreate)


Test patch:
diff --git a/pkg/metrics/metrics_test.go b/pkg/metrics/metrics_test.go
index fb58a62d..2bd1590d 100644
--- a/pkg/metrics/metrics_test.go
+++ b/pkg/metrics/metrics_test.go
@@ -16,6 +16,7 @@ func TestBasicAndDurationReport(t *testing.T) {
 
 	testCounterMetric(t, reporter, AssignedIdentityAdditionCountM)
 	testCounterMetric(t, reporter, AssignedIdentityDeletionCountM)
+	testCounterMetric(t, reporter, AssignedIdentityUpdateCountM)
 	testCounterMetric(t, reporter, MICCycleCountM)
 	testCounterMetric(t, reporter, MICNewLeaderElectionCountM)
 	testCounterMetric(t, reporter, CloudProviderOperationsErrorsCountM)
diff --git a/pkg/mic/mic_test.go b/pkg/mic/mic_test.go
index b5f2d3b8..5706d108 100644
--- a/pkg/mic/mic_test.go
+++ b/pkg/mic/mic_test.go
@@ -89,7 +89,6 @@ func (c *TestVMClient) UpdateIdentities(rg, nodeName string, vm compute.VirtualM
 	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 {
@@ -104,7 +103,6 @@ func (c *TestVMClient) UpdateIdentities(rg, nodeName string, vm compute.VirtualM
 			delete(c.nodeIDs[nodeName], k)
 		}
 	}
-
 	c.nodeMap[nodeName] = &vm
 	return nil
 }
@@ -457,6 +455,14 @@ func (c *TestCrdClient) CreateAssignedIdentity(assignedIdentity *internalaadpodi
 	return nil
 }
 
+func (c *TestCrdClient) UpdateAssignedIdentity(assignedIdentity *internalaadpodid.AzureAssignedIdentity) error {
+	assignedIdentityToStore := *assignedIdentity //Make a copy to store in the map.
+	c.mu.Lock()
+	c.assignedIDMap[assignedIdentity.Name] = &assignedIdentityToStore
+	c.mu.Unlock()
+	return nil
+}
+
 func (c *TestCrdClient) UpdateAzureAssignedIdentityStatus(assignedIdentity *internalaadpodid.AzureAssignedIdentity, status string) error {
 	assignedIdentity.Status.Status = status
 	assignedIdentityToStore := *assignedIdentity //Make a copy to store in the map.
@@ -909,15 +915,6 @@ func TestSimpleMICClient(t *testing.T) {
 		t.Fatalf("Assigned id not deleted")
 	}
 
-	/*
-		testPass = evtRecorder.Validate(&LastEvent{Type: "Normal", Reason: "binding removed",
-			Message: "Binding testbinding removed from node test-node for pod test-pod"})
-
-		if !testPass {
-			t.Errorf("event mismatch")
-		}
-	*/
-
 	// Test3: Error from cloud provider event test
 	err = errors.New("error returned from cloud provider")
 	cloudClient.SetError(err)
@@ -937,14 +934,6 @@ func TestSimpleMICClient(t *testing.T) {
 		t.Fatalf("expected status to be %s, got: %s", aadpodid.AssignedIDCreated, (*listAssignedIDs)[0].Status.Status)
 	}
 
-	/*
-		testPass = evtRecorder.Validate(&LastEvent{Type: "Warning", Reason: "binding apply error",
-			Message: "Applying binding testbinding node test-node for pod test-pod-default-test-id resulted in error error returned from cloud provider"})
-
-		if !testPass {
-			t.Errorf("event mismatch")
-		} */
-
 	// Test4: Removal error event test
 	// Reset the state to add the id.
 	cloudClient.UnSetError()
@@ -956,15 +945,158 @@ func TestSimpleMICClient(t *testing.T) {
 	cloudClient.SetError(err)
 
 	podClient.DeletePod("test-pod", "default")
+}
+
+func TestUpdateAssignedIdentities(t *testing.T) {
+	eventCh := make(chan internalaadpodid.EventType, 100)
+	cloudClient := NewTestCloudClient(config.AzureConfig{})
+	crdClient := NewTestCrdClient(nil)
+	podClient := NewTestPodClient()
+	nodeClient := NewTestNodeClient()
+	var evtRecorder TestEventRecorder
+	evtRecorder.lastEvent = new(LastEvent)
+	evtRecorder.eventChannel = make(chan bool, 100)
+
+	micClient := NewMICTestClient(eventCh, cloudClient, crdClient, podClient, nodeClient, &evtRecorder, false, 4, nil)
+
+	crdClient.CreateID("test-id", "default", aadpodid.UserAssignedMSI, testResourceID, "test-user-msi-clientid", nil, "", "", "", "rv1")
+	crdClient.CreateBinding("testbinding", "default", "test-id", "test-select", "")
+
+	nodeClient.AddNode("test-node")
+	podClient.AddPod("test-pod", "default", "test-node", "test-select")
+
+	eventCh <- internalaadpodid.PodCreated
+
+	defer micClient.testRunSync()(t)
+
+	evtRecorder.WaitForEvents(1)
+
+	testPass := false
+	listAssignedIDs, err := crdClient.ListAssignedIDs()
+	if err != nil {
+		klog.Error(err)
+		t.Errorf("list assigned failed")
+	}
+
+	if listAssignedIDs != nil {
+		for _, assignedID := range *listAssignedIDs {
+			if assignedID.Spec.Pod == "test-pod" && assignedID.Spec.PodNamespace == "default" && assignedID.Spec.NodeName == "test-node" &&
+				assignedID.Spec.AzureBindingRef.Name == "testbinding" && assignedID.Spec.AzureIdentityRef.Name == "test-id" &&
+				assignedID.Spec.AzureIdentityRef.ResourceVersion == "rv1" && assignedID.Spec.AzureIdentityRef.Spec.ClientID == "test-user-msi-clientid" {
+				testPass = true
+				break
+			}
+		}
+	}
+
+	if !testPass {
+		t.Fatalf("assigned id mismatch")
+	}
+	newResourceID := testResourceID + "-new"
+	crdClient.CreateID("test-id", "default", aadpodid.UserAssignedMSI, newResourceID, "test-user-msi-clientid", nil, "", "", "", "changedrv2")
+	crdClient.CreateID("test-id-2", "default", aadpodid.UserAssignedMSI, testResourceID, "test-user-msi-clientid", nil, "", "", "", "rv2")
+	crdClient.CreateBinding("testbinding2", "default", "test-id-2", "test-select", "")
+
+	eventCh <- internalaadpodid.IdentityUpdated
+	eventCh <- internalaadpodid.IdentityCreated
+	eventCh <- internalaadpodid.BindingCreated
+
+	evtRecorder.WaitForEvents(1)
+
+	listAssignedIDs, err = crdClient.ListAssignedIDs()
+	if err != nil {
+		klog.Error(err)
+		t.Errorf("list assigned failed")
+	}
+
+	testPass = false
+	if listAssignedIDs != nil {
+		for _, assignedID := range *listAssignedIDs {
+			if assignedID.Spec.Pod == "test-pod" && assignedID.Spec.PodNamespace == "default" && assignedID.Spec.NodeName == "test-node" &&
+				assignedID.Spec.AzureBindingRef.Name == "testbinding" && assignedID.Spec.AzureIdentityRef.Name == "test-id" &&
+				assignedID.Spec.AzureIdentityRef.ResourceVersion == "changedrv2" && assignedID.Spec.AzureIdentityRef.Spec.ClientID == "test-user-msi-clientid" &&
+				assignedID.Spec.AzureIdentityRef.Spec.ResourceID == newResourceID {
+				testPass = true
+				break
+			}
+		}
+	}
+	if !testPass {
+		t.Fatalf("assigned id mismatch")
+	}
+}
+
+func TestAddUpdateDel(t *testing.T) {
+	eventCh := make(chan internalaadpodid.EventType, 100)
+	cloudClient := NewTestCloudClient(config.AzureConfig{})
+	crdClient := NewTestCrdClient(nil)
+	podClient := NewTestPodClient()
+	nodeClient := NewTestNodeClient()
+	var evtRecorder TestEventRecorder
+	evtRecorder.lastEvent = new(LastEvent)
+	evtRecorder.eventChannel = make(chan bool, 100)
+
+	micClient := NewMICTestClient(eventCh, cloudClient, crdClient, podClient, nodeClient, &evtRecorder, false, 4, nil)
+
+	crdClient.CreateID("test-id-0", "default", aadpodid.UserAssignedMSI, fmt.Sprintf("%s-%d", testResourceID, 0), "test-user-msi-clientid-0", nil, "", "", "", "rv-0")
+	crdClient.CreateBinding("testbinding-0", "default", "test-id-0", "test-select-0", "")
+
+	crdClient.CreateID("test-id-1", "default", aadpodid.UserAssignedMSI, fmt.Sprintf("%s-%d", testResourceID, 1), "test-user-msi-clientid-1", nil, "", "", "", "rv-1")
+	crdClient.CreateBinding("testbinding-1", "default", "test-id-1", "test-select-1", "")
+
+	crdClient.CreateID("test-id-2", "default", aadpodid.UserAssignedMSI, fmt.Sprintf("%s-%d", testResourceID, 2), "test-user-msi-clientid-2", nil, "", "", "", "rv-2")
+	crdClient.CreateBinding("testbinding-2", "default", "test-id-2", "test-select-2", "")
+
+	nodeClient.AddNode("test-node")
+	podClient.AddPod("test-pod-0", "default", "test-node", "test-select-0")
+	podClient.AddPod("test-pod-1", "default", "test-node", "test-select-1")
+	podClient.AddPod("test-pod-2", "default", "test-node", "test-select-2")
+
+	eventCh <- internalaadpodid.PodCreated
+	eventCh <- internalaadpodid.PodCreated
+	eventCh <- internalaadpodid.PodCreated
+
+	defer micClient.testRunSync()(t)
+
+	evtRecorder.WaitForEvents(3)
+	listAssignedIDs, err := crdClient.ListAssignedIDs()
+	if err != nil {
+		t.Fatalf("failed to list assigned ids, err: %v", err)
+	}
+	if len(*listAssignedIDs) != 3 {
+		t.Fatalf("expected len to be %d, got %d", 3, len(*listAssignedIDs))
+	}
+
+	crdClient.CreateID("test-id-0", "default", aadpodid.UserAssignedMSI, fmt.Sprintf("%s-%d", testResourceID, 4), "test-user-msi-clientid-4", nil, "", "", "", "updated-rv-0")
+	crdClient.CreateID("test-id-3", "default", aadpodid.UserAssignedMSI, fmt.Sprintf("%s-%d", testResourceID, 3), "test-user-msi-clientid-3", nil, "", "", "", "rv-3")
+	crdClient.CreateBinding("testbinding-3", "default", "test-id-3", "test-select-2", "")
+	podClient.DeletePod("test-pod-1", "default")
+
+	eventCh <- internalaadpodid.IdentityCreated
+	eventCh <- internalaadpodid.BindingCreated
+	eventCh <- internalaadpodid.IdentityUpdated
 	eventCh <- internalaadpodid.PodDeleted
-	/*
-		testPass = evtRecorder.Validate(&LastEvent{Type: "Warning", Reason: "binding remove error",
-			Message: "Binding testbinding removal from node test-node for pod test-pod resulted in error remove error returned from cloud provider"})
 
-		if !testPass {
-			t.Errorf("event mismatch")
+	evtRecorder.WaitForEvents(3)
+	listAssignedIDs, err = crdClient.ListAssignedIDs()
+	if err != nil {
+		t.Fatalf("failed to list assigned ids, err: %v", err)
+	}
+	if len(*listAssignedIDs) != 3 {
+		t.Fatalf("expected len to be %d, got %d", 3, len(*listAssignedIDs))
+	}
+	// check the updated identity has the correct azureid ref
+	for _, assignedID := range *listAssignedIDs {
+		if assignedID.Name != "test-pod-0-default-test-id-0" {
+			continue
 		}
-	*/
+		if !(assignedID.Spec.Pod == "test-pod-0" && assignedID.Spec.PodNamespace == "default" && assignedID.Spec.NodeName == "test-node" &&
+			assignedID.Spec.AzureBindingRef.Name == "testbinding-0" && assignedID.Spec.AzureIdentityRef.Name == "test-id-0" &&
+			assignedID.Spec.AzureIdentityRef.ResourceVersion == "updated-rv-0" && assignedID.Spec.AzureIdentityRef.Spec.ClientID == "test-user-msi-clientid-4" &&
+			assignedID.Spec.AzureIdentityRef.Spec.ResourceID == fmt.Sprintf("%s-%d", testResourceID, 4)) {
+			t.Fatalf("azure identity spec mismatch")
+		}
+	}
 }
 
 func TestAddDelMICClient(t *testing.T) {
@@ -1206,7 +1338,7 @@ func TestMICStateFlow(t *testing.T) {
 	cloudClient.SetError(errors.New("error removing identity from node"))
 	cloudClient.testVMClient.identity = &compute.VirtualMachineIdentity{
 		UserAssignedIdentities: map[string]*compute.VirtualMachineIdentityUserAssignedIdentitiesValue{
-			testResourceID: &compute.VirtualMachineIdentityUserAssignedIdentitiesValue{},
+			testResourceID: {},
 		},
 	}
 
@@ -1391,7 +1523,7 @@ func TestSyncRetryLoop(t *testing.T) {
 	cloudClient.SetError(errors.New("error removing identity from node"))
 	cloudClient.testVMClient.identity = &compute.VirtualMachineIdentity{
 		UserAssignedIdentities: map[string]*compute.VirtualMachineIdentityUserAssignedIdentitiesValue{
-			testResourceID: &compute.VirtualMachineIdentityUserAssignedIdentitiesValue{},
+			testResourceID: {},
 		},
 	}
 
@@ -1568,7 +1700,6 @@ func TestSyncExit(t *testing.T) {
 
 	micClient.testRunSync()(t)
 }
-
 func TestMicAddDelVMSSwithImmutableIdentities(t *testing.T) {
 	eventCh := make(chan internalaadpodid.EventType, 100)
 	cloudClient := NewTestCloudClient(config.AzureConfig{VMType: "vmss"})
diff --git a/test/common/k8s/azureidentity/azureidentity.go b/test/common/k8s/azureidentity/azureidentity.go
index a95f06db..0de7c360 100644
--- a/test/common/k8s/azureidentity/azureidentity.go
+++ b/test/common/k8s/azureidentity/azureidentity.go
@@ -18,17 +18,17 @@ import (
 
 // CreateOnClusterOld will create an Azure Identity on a Kubernetes cluster
 func CreateOnClusterOld(subscriptionID, resourceGroup, name, templateOutputPath string) error {
-	return CreateOnClusterInternal(subscriptionID, resourceGroup, name, "aadpodidentity-old.yaml", templateOutputPath)
+	return CreateOnClusterInternal(subscriptionID, resourceGroup, name, name, "aadpodidentity-old.yaml", templateOutputPath)
 }
 
 // CreateOnCluster will create an Azure Identity on a Kubernetes cluster
-func CreateOnCluster(subscriptionID, resourceGroup, name, templateOutputPath string) error {
-	return CreateOnClusterInternal(subscriptionID, resourceGroup, name, "aadpodidentity.yaml", templateOutputPath)
+func CreateOnCluster(subscriptionID, resourceGroup, identity, name, templateOutputPath string) error {
+	return CreateOnClusterInternal(subscriptionID, resourceGroup, identity, name, "aadpodidentity.yaml", templateOutputPath)
 }
 
 // CreateOnClusterInternal will create an Azure Identity on a Kubernetes cluster
-func CreateOnClusterInternal(subscriptionID, resourceGroup, name, templateInputFile, templateOutputPath string) error {
-	clientID, err := GetClientID(resourceGroup, name)
+func CreateOnClusterInternal(subscriptionID, resourceGroup, identity, name, templateInputFile, templateOutputPath string) error {
+	clientID, err := GetClientID(resourceGroup, identity)
 	if err != nil {
 		return err
 	}
@@ -51,11 +51,13 @@ func CreateOnClusterInternal(subscriptionID, resourceGroup, name, templateInputF
 		ResourceGroup  string
 		ClientID       string
 		Name           string
+		Identity       string
 	}{
 		subscriptionID,
 		resourceGroup,
 		clientID,
 		name,
+		identity,
 	}
 	if err := t.Execute(deployFile, deployData); err != nil {
 		return errors.Wrap(err, "Failed to create a deployment file from aadpodidentity.yaml")
diff --git a/test/e2e/aadpodidentity_test.go b/test/e2e/aadpodidentity_test.go
index 015f4a58..72802d0c 100644
--- a/test/e2e/aadpodidentity_test.go
+++ b/test/e2e/aadpodidentity_test.go
@@ -146,7 +146,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 		azureAssignedIdentity, err := azureassignedidentity.GetByPrefix(identityValidator)
 		Expect(err).NotTo(HaveOccurred())
 
-		validateAzureAssignedIdentity(azureAssignedIdentity, keyvaultIdentity)
+		validateAzureAssignedIdentity(azureAssignedIdentity, keyvaultIdentity, keyvaultIdentity)
 	})
 
 	It("should not pass the identity validating test if the AzureIdentity is deleted", func() {
@@ -226,7 +226,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 		Expect(err).NotTo(HaveOccurred())
 
 		// Make sure the AzureAssignedIdentity is updated along with the new pod
-		validateAzureAssignedIdentity(azureAssignedIdentity, keyvaultIdentity)
+		validateAzureAssignedIdentity(azureAssignedIdentity, keyvaultIdentity, keyvaultIdentity)
 
 		node.Uncordon(nodeName)
 	})
@@ -256,7 +256,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 			azureAssignedIdentity, err := azureassignedidentity.GetByPrefix(identityValidatorName)
 			Expect(err).NotTo(HaveOccurred())
 
-			validateAzureAssignedIdentity(azureAssignedIdentity, identityName)
+			validateAzureAssignedIdentity(azureAssignedIdentity, identityName, identityName)
 
 			testData[i] = struct {
 				identityName          string
@@ -293,7 +293,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 
 			// Make sure that the existing identities are still functioning
 			for j := i + 1; j < 5; j++ {
-				validateAzureAssignedIdentity(testData[j].azureAssignedIdentity, testData[j].identityName)
+				validateAzureAssignedIdentity(testData[j].azureAssignedIdentity, testData[j].identityName, testData[j].identityName)
 			}
 		}
 	})
@@ -342,7 +342,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 
 		azureAssignedIdentity, err := azureassignedidentity.GetByPrefix(identityValidator)
 		Expect(err).NotTo(HaveOccurred())
-		validateAzureAssignedIdentity(azureAssignedIdentity, keyvaultIdentity)
+		validateAzureAssignedIdentity(azureAssignedIdentity, keyvaultIdentity, keyvaultIdentity)
 
 		// Delete pod identity to verify that the VM identity did not get deleted
 		waitForDeployDeletion(identityValidator)
@@ -430,7 +430,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 
 		azureAssignedIdentity, err := azureassignedidentity.GetByPrefix(identityValidator)
 		Expect(err).NotTo(HaveOccurred())
-		validateAzureAssignedIdentity(azureAssignedIdentity, keyvaultIdentity)
+		validateAzureAssignedIdentity(azureAssignedIdentity, keyvaultIdentity, keyvaultIdentity)
 
 		waitForDeployDeletion(identityValidator)
 		ok, err := azureassignedidentity.WaitOnLengthMatched(0, false)
@@ -472,7 +472,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 			Expect(err).NotTo(HaveOccurred())
 
 			for _, azureAssignedIdentity := range azureAssignedIdentities {
-				validateAzureAssignedIdentity(azureAssignedIdentity, identityName)
+				validateAzureAssignedIdentity(azureAssignedIdentity, identityName, identityName)
 			}
 		}
 	})
@@ -523,7 +523,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 		azureAssignedIdentity, err := azureassignedidentity.GetByPrefix(identityValidator)
 		Expect(err).NotTo(HaveOccurred())
 
-		validateAzureAssignedIdentity(azureAssignedIdentity, keyvaultIdentity)
+		validateAzureAssignedIdentity(azureAssignedIdentity, keyvaultIdentity, keyvaultIdentity)
 	})
 
 	It("should not delete an in use identity from a vmss", func() {
@@ -603,7 +603,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 		Expect(err).NotTo(HaveOccurred())
 
 		// Ensure that the identity validator is able to get the token
-		validateAzureAssignedIdentity(azureAssignedIdentity, immutableIdentity)
+		validateAzureAssignedIdentity(azureAssignedIdentity, immutableIdentity, immutableIdentity)
 
 		waitForDeployDeletion(identityValidator)
 
@@ -753,7 +753,8 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 
 		setUpIdentityAndDeployment(keyvaultIdentity, "", "1")
 
-		err := azureidentity.CreateOnCluster(cfg.SubscriptionID, cfg.ResourceGroup, fmt.Sprintf("%s-%d", keyvaultIdentity, 5), templateOutputPath)
+		azureIdentity := fmt.Sprintf("%s-%d", keyvaultIdentity, 5)
+		err := azureidentity.CreateOnCluster(cfg.SubscriptionID, cfg.ResourceGroup, azureIdentity, azureIdentity, templateOutputPath)
 		Expect(err).NotTo(HaveOccurred())
 
 		err = azureidentitybinding.Create(fmt.Sprintf("%s-%d", keyvaultIdentity, 5), keyvaultIdentity, templateOutputPath)
@@ -790,7 +791,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 		azureAssignedIdentity, err := azureassignedidentity.GetByPrefix(identityValidator)
 		Expect(err).NotTo(HaveOccurred())
 
-		validateAzureAssignedIdentity(azureAssignedIdentity, keyvaultIdentity)
+		validateAzureAssignedIdentity(azureAssignedIdentity, keyvaultIdentity, keyvaultIdentity)
 
 		nodeList, err := node.GetAll()
 		Expect(err).NotTo(HaveOccurred())
@@ -842,11 +843,10 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 		azureAssignedIdentity, err := azureassignedidentity.GetByPrefix(identityValidator)
 		Expect(err).NotTo(HaveOccurred())
 
-		validateAzureAssignedIdentity(azureAssignedIdentity, keyvaultIdentity)
+		validateAzureAssignedIdentity(azureAssignedIdentity, keyvaultIdentity, keyvaultIdentity)
 	})
 
 	It("should pass the identity format validation with gatekeeper constraint", func() {
-
 		// setup the required infra
 		setupIdentityFormatValidationInfra()
 
@@ -867,6 +867,31 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 		fmt.Printf("%s", output)
 		Expect(err).NotTo(HaveOccurred())
 	})
+
+	It("should assign new identity and remove old when AzureIdentity updated", func() {
+		setUpIdentityAndDeployment(keyvaultIdentity, "", "1")
+		ok, err := azureassignedidentity.WaitOnLengthMatched(1, true)
+		Expect(err).NotTo(HaveOccurred())
+		Expect(ok).To(Equal(true))
+
+		azureAssignedIdentity, err := azureassignedidentity.GetByPrefix(identityValidator)
+		Expect(err).NotTo(HaveOccurred())
+
+		validateAzureAssignedIdentity(azureAssignedIdentity, keyvaultIdentity, keyvaultIdentity)
+
+		identity := fmt.Sprintf("%s-%d", keyvaultIdentity, 0)
+		err = azureidentity.CreateOnCluster(cfg.SubscriptionID, cfg.ResourceGroup, identity, keyvaultIdentity, templateOutputPath)
+		Expect(err).NotTo(HaveOccurred())
+
+		ok, err = azureassignedidentity.WaitOnLengthMatched(1, true)
+		Expect(err).NotTo(HaveOccurred())
+		Expect(ok).To(Equal(true))
+
+		azureAssignedIdentity, err = azureassignedidentity.GetByPrefix(identityValidator)
+		Expect(err).NotTo(HaveOccurred())
+
+		validateAzureAssignedIdentity(azureAssignedIdentity, keyvaultIdentity, identity)
+	})
 })
 
 func GetVMSS(nodeList *node.List) ([]node.Node, string) {
@@ -925,7 +950,7 @@ func runValidatorTest(iterations int) {
 
 		azureAssignedIdentity, err := azureassignedidentity.GetByPrefix(identityValidator)
 		Expect(err).NotTo(HaveOccurred())
-		validateAzureAssignedIdentity(azureAssignedIdentity, keyvaultIdentity)
+		validateAzureAssignedIdentity(azureAssignedIdentity, keyvaultIdentity, keyvaultIdentity)
 
 		deleteAllIdentityValidator()
 
@@ -1151,7 +1176,7 @@ func setUpIdentityAndDeploymentInternal(azureIdentityName, suffix, replicas stri
 	if old {
 		err = azureidentity.CreateOnClusterOld(cfg.SubscriptionID, cfg.ResourceGroup, azureIdentityName, templateOutputPath)
 	} else {
-		err = azureidentity.CreateOnCluster(cfg.SubscriptionID, cfg.ResourceGroup, azureIdentityName, templateOutputPath)
+		err = azureidentity.CreateOnCluster(cfg.SubscriptionID, cfg.ResourceGroup, azureIdentityName, azureIdentityName, templateOutputPath)
 	}
 	Expect(err).NotTo(HaveOccurred())
 
@@ -1190,7 +1215,7 @@ func setUpIdentityAndDeploymentInternal(azureIdentityName, suffix, replicas stri
 }
 
 // validateAzureAssignedIdentity will make sure a given AzureAssignedIdentity has the correct properties
-func validateAzureAssignedIdentity(azureAssignedIdentity aadpodid.AzureAssignedIdentity, identityName string) {
+func validateAzureAssignedIdentity(azureAssignedIdentity aadpodid.AzureAssignedIdentity, identityName, identity string) {
 	identityClientID := azureAssignedIdentity.Spec.AzureIdentityRef.Spec.ClientID
 	Expect(identityClientID).NotTo(Equal(""))
 	podName := azureAssignedIdentity.Spec.Pod
@@ -1209,7 +1234,7 @@ func validateAzureAssignedIdentity(azureAssignedIdentity aadpodid.AzureAssignedI
 	Expect(azureAssignedIdentity.Spec.AzureIdentityRef.ObjectMeta.Name).To(Equal(identityName))
 	Expect(azureAssignedIdentity.Spec.AzureIdentityRef.ObjectMeta.Namespace).To(Equal("default"))
 
-	if strings.HasPrefix(identityName, clusterIdentity) {
+	if strings.HasPrefix(identity, clusterIdentity) {
 		cmdOutput, err := validateClusterWideUserAssignedIdentity(podName, identityClientID)
 		Expect(errors.Wrap(err, string(cmdOutput))).NotTo(HaveOccurred())
 	} else {
@@ -1218,7 +1243,7 @@ func validateAzureAssignedIdentity(azureAssignedIdentity aadpodid.AzureAssignedI
 		Expect(errors.Wrap(err, string(cmdOutput))).NotTo(HaveOccurred())
 	}
 
-	fmt.Printf("# %s validated!\n", identityName)
+	fmt.Printf("# %s validated!\n", identity)
 }
 
 // waitForDeployDeletion will block until a give deploy and its pods are completed deleted
diff --git a/test/e2e/template/aadpodidentity.yaml b/test/e2e/template/aadpodidentity.yaml
index 274c776e..1a5ef6e9 100644
--- a/test/e2e/template/aadpodidentity.yaml
+++ b/test/e2e/template/aadpodidentity.yaml
@@ -4,5 +4,5 @@ metadata:
   name: {{.Name}}
 spec:
   type: 0
-  resourceID: /subscriptions/{{.SubscriptionID}}/resourceGroups/{{.ResourceGroup}}/providers/Microsoft.ManagedIdentity/userAssignedIdentities/{{.Name}}
+  resourceID: /subscriptions/{{.SubscriptionID}}/resourceGroups/{{.ResourceGroup}}/providers/Microsoft.ManagedIdentity/userAssignedIdentities/{{.Identity}}
   clientID: {{.ClientID}}