Repository: Azure/aad-pod-identity
Pull request #486 (2020-01-16 18:35:38)

Interface phase2 - implement managed mode
<!-- 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? -->


**Issue Fixed**:
<!-- If this PR fixes GitHub issue 1234, add "Fixes #1234" to the next line. -->

**Notes for Reviewers**:


aramase: ```
Ran 19 of 20 Specs in 3549.162 seconds
SUCCESS! -- 19 Passed | 0 Failed | 0 Pending | 1 Skipped
--- PASS: TestAADPodIdentity (3549.16s)
PASS
ok	github.com/Azure/aad-pod-identity/test/e2e	3550.169s
```
kkmsft: Are we planning to have separate PR for helm chart changes ? (would really prefer that :-) )

Patch:
diff --git a/cmd/nmi/main.go b/cmd/nmi/main.go
index f8948b97..ef460894 100644
--- a/cmd/nmi/main.go
+++ b/cmd/nmi/main.go
@@ -3,12 +3,13 @@ package main
 import (
 	goflag "flag"
 	"os"
+	"strings"
 
 	"net/http"
 	_ "net/http/pprof"
 
-	"github.com/Azure/aad-pod-identity/pkg/k8s"
 	"github.com/Azure/aad-pod-identity/pkg/metrics"
+	"github.com/Azure/aad-pod-identity/pkg/nmi"
 	server "github.com/Azure/aad-pod-identity/pkg/nmi/server"
 	"github.com/Azure/aad-pod-identity/pkg/probes"
 	"github.com/Azure/aad-pod-identity/version"
@@ -45,10 +46,10 @@ var (
 	enableScaleFeatures                = pflag.Bool("enableScaleFeatures", false, "Enable/Disable features for scale clusters")
 	blockInstanceMetadata              = pflag.Bool("block-instance-metadata", false, "Block instance metadata endpoints")
 	prometheusPort                     = pflag.String("prometheus-port", "9090", "Prometheus port for metrics")
+	operationMode                      = pflag.String("operation-mode", "standard", "NMI operation mode")
 )
 
 func main() {
-	klog.InitFlags(nil)
 	// this is done for glog used by client-go underneath
 	pflag.CommandLine.AddGoFlagSet(goflag.CommandLine)
 
@@ -70,14 +71,18 @@ func main() {
 		klog.Infof("Features for scale clusters enabled")
 	}
 
-	client, err := k8s.NewKubeClient(*nodename, *enableScaleFeatures)
+	// normalize operation mode
+	*operationMode = strings.ToLower(*operationMode)
+
+	client, err := nmi.GetKubeClient(*nodename, *operationMode, *enableScaleFeatures)
 	if err != nil {
-		klog.Fatalf("%+v", err)
+		klog.Fatalf("error creating kube client, err: %+v", err)
 	}
+
 	exit := make(<-chan struct{})
 	client.Start(exit)
 	*forceNamespaced = *forceNamespaced || "true" == os.Getenv("FORCENAMESPACED")
-	s := server.NewServer(*forceNamespaced, *micNamespace, *blockInstanceMetadata)
+	s := server.NewServer(*micNamespace, *blockInstanceMetadata)
 	s.KubeClient = client
 	s.MetadataIP = *metadataIP
 	s.MetadataPort = *metadataPort
@@ -85,9 +90,21 @@ func main() {
 	s.HostIP = *hostIP
 	s.NodeName = *nodename
 	s.IPTableUpdateTimeIntervalInSeconds = *ipTableUpdateTimeIntervalInSeconds
-	s.ListPodIDsRetryAttemptsForCreated = *retryAttemptsForCreated
-	s.ListPodIDsRetryAttemptsForAssigned = *retryAttemptsForAssigned
-	s.ListPodIDsRetryIntervalInSeconds = *findIdentityRetryIntervalInSeconds
+
+	nmiConfig := nmi.Config{
+		Mode:                               strings.ToLower(*operationMode),
+		RetryAttemptsForCreated:            *retryAttemptsForCreated,
+		RetryAttemptsForAssigned:           *retryAttemptsForAssigned,
+		FindIdentityRetryIntervalInSeconds: *findIdentityRetryIntervalInSeconds,
+		Namespaced:                         *forceNamespaced,
+	}
+
+	// Create new token client based on the nmi mode
+	tokenClient, err := nmi.GetTokenClient(client, nmiConfig)
+	if err != nil {
+		klog.Fatalf("failed to initialize token client, err: %v", err)
+	}
+	s.TokenClient = tokenClient
 
 	// Health probe will always report success once its started. The contents
 	// will report "Active" once the iptables rules are set
@@ -97,7 +114,6 @@ func main() {
 	if err = metrics.RegisterAndExport(*prometheusPort); err != nil {
 		klog.Fatalf("Could not register and export metrics: %+v", err)
 	}
-
 	if err := s.Run(); err != nil {
 		klog.Fatalf("%s", err)
 	}
diff --git a/go.mod b/go.mod
index dcd1a018..9fdc8de9 100644
--- a/go.mod
+++ b/go.mod
@@ -14,6 +14,7 @@ require (
 	github.com/coreos/go-iptables v0.3.0
 	github.com/gogo/protobuf v1.2.1 // indirect
 	github.com/golang/groupcache v0.0.0-20180513044358-24b0969c4cb7 // indirect
+	github.com/google/go-cmp v0.3.0
 	github.com/google/gofuzz v0.0.0-20170612174753-24818f796faf // indirect
 	github.com/googleapis/gnostic v0.1.0 // indirect
 	github.com/howeyc/gopass v0.0.0-20170109162249-bf9dde6d0d2c // indirect
diff --git a/pkg/crd/crd.go b/pkg/crd/crd.go
index ad07ff63..499a912d 100644
--- a/pkg/crd/crd.go
+++ b/pkg/crd/crd.go
@@ -35,8 +35,7 @@ type Client struct {
 // ClientInt ...
 type ClientInt interface {
 	Start(exit <-chan struct{})
-	SyncCache(exit <-chan struct{})
-	SyncCacheLite(exit <-chan struct{})
+	SyncCache(exit <-chan struct{}, initial bool, cacheSyncs ...cache.InformerSynced)
 	RemoveAssignedIdentity(assignedIdentity *aadpodid.AzureAssignedIdentity) error
 	CreateAssignedIdentity(assignedIdentity *aadpodid.AzureAssignedIdentity) error
 	UpdateAzureAssignedIdentityStatus(assignedIdentity *aadpodid.AzureAssignedIdentity, status string) error
@@ -49,25 +48,37 @@ type ClientInt interface {
 }
 
 // NewCRDClientLite ...
-func NewCRDClientLite(config *rest.Config, nodeName string, scale bool) (crdClient *Client, err error) {
+func NewCRDClientLite(config *rest.Config, nodeName string, scale, isStandardMode bool) (crdClient *Client, err error) {
 	restClient, err := newRestClient(config)
 	if err != nil {
 		klog.Error(err)
 		return nil, err
 	}
 
-	var assignedIDListWatch *cache.ListWatch
+	var assignedIDListInformer, bindingListInformer, idListInformer cache.SharedInformer
 
-	if scale {
-		assignedIDListWatch = newAssignedIDNodeListWatch(restClient, nodeName)
-	} else {
-		assignedIDListWatch = newAssignedIDListWatch(restClient)
-	}
+	// assigned identity informer is required only for standard mode
+	if isStandardMode {
+		var assignedIDListWatch *cache.ListWatch
+		if scale {
+			assignedIDListWatch = newAssignedIDNodeListWatch(restClient, nodeName)
+		} else {
+			assignedIDListWatch = newAssignedIDListWatch(restClient)
+		}
 
-	assignedIDListInformer, err := newAssignedIDInformer(assignedIDListWatch)
-	if err != nil {
-		klog.Error(err)
-		return nil, err
+		assignedIDListInformer, err = newAssignedIDInformer(assignedIDListWatch)
+		if err != nil {
+			klog.Error(err)
+			return nil, err
+		}
+	} else {
+		// creating binding and identity list informers for non standard mode
+		if bindingListInformer, err = newBindingInformerLite(newBindingListWatch(restClient)); err != nil {
+			return nil, err
+		}
+		if idListInformer, err = newIDInformerLite(newIDListWatch(restClient)); err != nil {
+			return nil, err
+		}
 	}
 	podIdentityExceptionListWatch := newPodIdentityExceptionListWatch(restClient)
 	podIdentityExceptionInformer, err := newPodIdentityExceptionInformer(podIdentityExceptionListWatch)
@@ -85,6 +96,8 @@ func NewCRDClientLite(config *rest.Config, nodeName string, scale bool) (crdClie
 	return &Client{
 		AssignedIDInformer:           assignedIDListInformer,
 		PodIdentityExceptionInformer: podIdentityExceptionInformer,
+		BindingInformer:              bindingListInformer,
+		IDInformer:                   idListInformer,
 		rest:                         restClient,
 		reporter:                     reporter,
 	}, nil
@@ -249,10 +262,25 @@ func newAssignedIDInformer(lw *cache.ListWatch) (cache.SharedInformer, error) {
 	if azAssignedIDInformer == nil {
 		return nil, fmt.Errorf("could not create %s informer", aadpodv1.AzureAssignedIDResource)
 	}
-
 	return azAssignedIDInformer, nil
 }
 
+func newBindingInformerLite(lw *cache.ListWatch) (cache.SharedInformer, error) {
+	azBindingInformer := cache.NewSharedInformer(lw, &aadpodv1.AzureIdentityBinding{}, time.Minute*10)
+	if azBindingInformer == nil {
+		return nil, fmt.Errorf("could not create %s informer", aadpodv1.AzureIDBindingResource)
+	}
+	return azBindingInformer, nil
+}
+
+func newIDInformerLite(lw *cache.ListWatch) (cache.SharedInformer, error) {
+	azIDInformer := cache.NewSharedInformer(lw, &aadpodv1.AzureIdentity{}, time.Minute*10)
+	if azIDInformer == nil {
+		return nil, fmt.Errorf("could not create %s informer", aadpodv1.AzureIDResource)
+	}
+	return azIDInformer, nil
+}
+
 func newPodIdentityExceptionListWatch(r *rest.RESTClient) *cache.ListWatch {
 	optionsModifier := func(options *v1.ListOptions) {}
 	return cache.NewFilteredListWatchFromClient(
@@ -273,9 +301,25 @@ func newPodIdentityExceptionInformer(lw *cache.ListWatch) (cache.SharedInformer,
 
 // StartLite to be used only case of lite client
 func (c *Client) StartLite(exit <-chan struct{}) {
-	go c.AssignedIDInformer.Run(exit)
-	go c.PodIdentityExceptionInformer.Run(exit)
-	c.SyncCacheLite(exit)
+	var cacheHasSynced []cache.InformerSynced
+
+	if c.AssignedIDInformer != nil {
+		go c.AssignedIDInformer.Run(exit)
+		cacheHasSynced = append(cacheHasSynced, c.AssignedIDInformer.HasSynced)
+	}
+	if c.BindingInformer != nil {
+		go c.BindingInformer.Run(exit)
+		cacheHasSynced = append(cacheHasSynced, c.BindingInformer.HasSynced)
+	}
+	if c.IDInformer != nil {
+		go c.IDInformer.Run(exit)
+		cacheHasSynced = append(cacheHasSynced, c.IDInformer.HasSynced)
+	}
+	if c.PodIdentityExceptionInformer != nil {
+		go c.PodIdentityExceptionInformer.Run(exit)
+		cacheHasSynced = append(cacheHasSynced, c.PodIdentityExceptionInformer.HasSynced)
+	}
+	c.SyncCache(exit, true, cacheHasSynced...)
 	klog.Info("CRD lite informers started ")
 }
 
@@ -284,23 +328,12 @@ func (c *Client) Start(exit <-chan struct{}) {
 	go c.BindingInformer.Run(exit)
 	go c.IDInformer.Run(exit)
 	go c.AssignedIDInformer.Run(exit)
-	c.SyncCache(exit)
+	c.SyncCache(exit, true, c.BindingInformer.HasSynced, c.IDInformer.HasSynced, c.AssignedIDInformer.HasSynced)
 	klog.Info("CRD informers started")
 }
 
-func (c *Client) SyncCache(exit <-chan struct{}) {
-	c.syncCache(exit, true, c.BindingInformer.HasSynced,
-		c.IDInformer.HasSynced,
-		c.AssignedIDInformer.HasSynced)
-}
-
-func (c *Client) SyncCacheLite(exit <-chan struct{}) {
-	c.syncCache(exit, true, c.AssignedIDInformer.HasSynced,
-		c.PodIdentityExceptionInformer.HasSynced)
-}
-
 // SyncCache synchronizes cache
-func (c *Client) syncCache(exit <-chan struct{}, initial bool, cacheSyncs ...cache.InformerSynced) {
+func (c *Client) SyncCache(exit <-chan struct{}, initial bool, cacheSyncs ...cache.InformerSynced) {
 	if !cache.WaitForCacheSync(exit, cacheSyncs...) {
 		if !initial {
 			klog.Errorf("Cache could not be synchronized")
@@ -536,6 +569,43 @@ func (c *Client) ListPodIds(podns, podname string) (map[string][]aadpodid.AzureI
 	return idStateMap, nil
 }
 
+// GetPodIDsWithBinding returns list of azure identity based on bindings
+// that match pod label.
+func (c *Client) GetPodIDsWithBinding(namespace string, labels map[string]string) ([]aadpodid.AzureIdentity, error) {
+	// get all bindings
+	bindings, err := c.ListBindings()
+	if err != nil {
+		return nil, err
+	}
+	if bindings == nil {
+		return nil, fmt.Errorf("binding list is nil from cache")
+	}
+	matchingIds := make(map[string]bool)
+	podLabel := labels[aadpodid.CRDLabelKey]
+
+	for _, binding := range *bindings {
+		// check if binding selector in pod labels
+		if podLabel == binding.Spec.Selector && binding.Namespace == namespace {
+			matchingIds[binding.Spec.AzureIdentity] = true
+		}
+	}
+	// get the azure identity objects based on the list generated
+	azIdentities, err := c.ListIds()
+	if err != nil {
+		return nil, err
+	}
+	if azIdentities == nil {
+		return nil, fmt.Errorf("azure identities list is nil from cache")
+	}
+	var azIds []aadpodid.AzureIdentity
+	for _, azIdentity := range *azIdentities {
+		if _, exists := matchingIds[azIdentity.Name]; exists && azIdentity.Namespace == namespace {
+			azIds = append(azIds, azIdentity)
+		}
+	}
+	return azIds, nil
+}
+
 type patchStatusOps struct {
 	Op    string      `json:"op"`
 	Path  string      `json:"path"`
diff --git a/pkg/k8s/client.go b/pkg/k8s/client.go
index bda1981e..a5196bca 100644
--- a/pkg/k8s/client.go
+++ b/pkg/k8s/client.go
@@ -32,10 +32,14 @@ const (
 type Client interface {
 	// Start just starts any informers required.
 	Start(<-chan struct{})
+	// GetPod returns the pod object based on name and namespce
+	GetPod(namespace, name string) (v1.Pod, error)
 	// GetPodInfo returns the pod name, namespace & replica set name for a given pod ip
 	GetPodInfo(podip string) (podns, podname, rsName string, selectors *metav1.LabelSelector, err error)
 	// ListPodIds pod matching azure identity or nil
 	ListPodIds(podns, podname string) (map[string][]aadpodid.AzureIdentity, error)
+	// ListPodIdsWithBinding pod matching azure identity or nil
+	ListPodIdsWithBinding(podns string, labels map[string]string) ([]aadpodid.AzureIdentity, error)
 	// GetSecret returns secret the secretRef represents
 	GetSecret(secretRef *v1.SecretReference) (*v1.Secret, error)
 	// ListPodIdentityExceptions returns list of azurepodidentityexceptions
@@ -53,7 +57,7 @@ type KubeClient struct {
 }
 
 // NewKubeClient new kubernetes api client
-func NewKubeClient(nodeName string, scale bool) (Client, error) {
+func NewKubeClient(nodeName string, scale, isStandardMode bool) (Client, error) {
 	config, err := buildConfig()
 	if err != nil {
 		return nil, err
@@ -63,7 +67,7 @@ func NewKubeClient(nodeName string, scale bool) (Client, error) {
 	if err != nil {
 		return nil, err
 	}
-	crdclient, err := crd.NewCRDClientLite(config, nodeName, scale)
+	crdclient, err := crd.NewCRDClientLite(config, nodeName, scale, isStandardMode)
 	if err != nil {
 		return nil, err
 	}
@@ -86,11 +90,11 @@ func NewKubeClient(nodeName string, scale bool) (Client, error) {
 	return kubeClient, nil
 }
 
+// Sync ...
 func (c *KubeClient) Sync(exit <-chan struct{}) {
 	if !cache.WaitForCacheSync(exit, c.PodInformer.HasSynced) {
 		klog.Errorf("Pod cache could not be synchronized")
 	}
-	c.CrdClient.SyncCacheLite(exit)
 }
 
 // Start the corresponding starts
@@ -121,6 +125,23 @@ func NodeNameFilter(nodeName string) internalinterfaces.TweakListOptionsFunc {
 	}
 }
 
+// GetPod returns pod that matches namespace and name
+func (c *KubeClient) GetPod(namespace, name string) (v1.Pod, error) {
+	// TODO (aramase) wrap this with retries
+	obj, exists, err := c.PodInformer.GetIndexer().GetByKey(namespace + "/" + name)
+	if err != nil {
+		return v1.Pod{}, fmt.Errorf("failed to get pod %s/%s, err: %+v", namespace, name, err)
+	}
+	if !exists {
+		return v1.Pod{}, fmt.Errorf("pod with key %s/%s doesn't exist", namespace, name)
+	}
+	pod, ok := obj.(*v1.Pod)
+	if !ok {
+		return v1.Pod{}, fmt.Errorf("could not cast %T to %s", pod, "v1.Pod")
+	}
+	return *pod, nil
+}
+
 // GetPodInfo get pod ns,name from apiserver
 func (c *KubeClient) GetPodInfo(podip string) (podns, poddname, rsName string, labels *metav1.LabelSelector, err error) {
 	if podip == "" {
@@ -214,6 +235,11 @@ func (c *KubeClient) ListPodIds(podns, podname string) (map[string][]aadpodid.Az
 	return c.CrdClient.ListPodIds(podns, podname)
 }
 
+// ListPodIdsWithBinding list matching ids for pod based on the bindings
+func (c *KubeClient) ListPodIdsWithBinding(podns string, labels map[string]string) ([]aadpodid.AzureIdentity, error) {
+	return c.CrdClient.GetPodIDsWithBinding(podns, labels)
+}
+
 // ListPodIdentityExceptions lists azurepodidentityexceptions
 func (c *KubeClient) ListPodIdentityExceptions(ns string) (*[]aadpodid.AzurePodIdentityException, error) {
 	return c.CrdClient.ListPodIdentityExceptions(ns)
diff --git a/pkg/k8s/fake.go b/pkg/k8s/fake.go
index a0021fd2..bb9f6e58 100644
--- a/pkg/k8s/fake.go
+++ b/pkg/k8s/fake.go
@@ -23,11 +23,21 @@ func (c *FakeClient) GetPodInfo(podip string) (podns, podname, rsName string, se
 	return "ns", "podname", "rsName", nil, nil
 }
 
+// GetPod returns fake pod object and nil error
+func (c *FakeClient) GetPod(podns, podname string) (v1.Pod, error) {
+	return v1.Pod{}, nil
+}
+
 // ListPodIds for pod
 func (c *FakeClient) ListPodIds(podns, podname string) (map[string][]aadpodid.AzureIdentity, error) {
 	return nil, nil
 }
 
+// ListPodIdsWithBinding for pod
+func (c *FakeClient) ListPodIdsWithBinding(podns string, labels map[string]string) ([]aadpodid.AzureIdentity, error) {
+	return nil, nil
+}
+
 // ListPodIdentityExceptions ...
 func (c *FakeClient) ListPodIdentityExceptions(ns string) (*[]aadpodid.AzurePodIdentityException, error) {
 	return nil, nil
diff --git a/pkg/nmi/managed.go b/pkg/nmi/managed.go
new file mode 100644
index 00000000..8be14ec0
--- /dev/null
+++ b/pkg/nmi/managed.go
@@ -0,0 +1,89 @@
+package nmi
+
+import (
+	"context"
+	"fmt"
+	"strings"
+
+	aadpodid "github.com/Azure/aad-pod-identity/pkg/apis/aadpodidentity"
+	auth "github.com/Azure/aad-pod-identity/pkg/auth"
+	k8s "github.com/Azure/aad-pod-identity/pkg/k8s"
+	utils "github.com/Azure/aad-pod-identity/pkg/utils"
+	"github.com/Azure/go-autorest/autorest/adal"
+	"k8s.io/klog"
+)
+
+// ManagedClient implements the TokenClient interface
+type ManagedClient struct {
+	TokenClient
+	KubeClient   k8s.Client
+	IsNamespaced bool
+}
+
+// NewManagedTokenClient creates new managed token client
+func NewManagedTokenClient(client k8s.Client, config Config) (*ManagedClient, error) {
+	// managed mode supported only in force namespaced mode
+	if !config.Namespaced {
+		return nil, fmt.Errorf("managed mode not intialized in force namespaced mode")
+	}
+	return &ManagedClient{
+		KubeClient:   client,
+		IsNamespaced: config.Namespaced,
+	}, nil
+}
+
+// GetIdentities gets the azure identity that matches the podns/podname and client id
+func (mc *ManagedClient) GetIdentities(ctx context.Context, podns, podname, clientID string) (*aadpodid.AzureIdentity, error) {
+	// get pod object to retrieve labels
+	pod, err := mc.KubeClient.GetPod(podns, podname)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get pod %s/%s, err: %+v", podns, podname, err)
+	}
+	// get all the azure identities based on azure identity bindings
+	azureIdentities, err := mc.KubeClient.ListPodIdsWithBinding(podns, pod.Labels)
+	for _, id := range azureIdentities {
+		// if client id exists in the request, then send the first identity that matched the client id
+		if len(clientID) != 0 && id.Spec.ClientID == clientID {
+			return &id, nil
+		}
+		// if client doesn't exist in the request, then return the first identity in the same namespace as the pod
+		if len(clientID) == 0 && strings.EqualFold(id.Namespace, podns) {
+			klog.Infof("No clientID in request. %s/%s has been matched with azure identity %s/%s", podns, podname, id.Namespace, id.Name)
+			return &id, nil
+		}
+	}
+	return nil, fmt.Errorf("no matching azure identity found for pod")
+}
+
+// GetToken ...
+func (mc *ManagedClient) GetToken(ctx context.Context, rqClientID, rqResource string, azureID aadpodid.AzureIdentity) (token *adal.Token, err error) {
+	rqHasClientID := len(rqClientID) != 0
+	clientID := azureID.Spec.ClientID
+
+	idType := azureID.Spec.Type
+	switch idType {
+	case aadpodid.UserAssignedMSI:
+		if rqHasClientID && !strings.EqualFold(rqClientID, clientID) {
+			klog.Warningf("clientid mismatch, requested:%s available:%s", rqClientID, clientID)
+		}
+		klog.Infof("matched identityType:%v clientid:%s resource:%s", idType, utils.RedactClientID(clientID), rqResource)
+		token, err := auth.GetServicePrincipalTokenFromMSIWithUserAssignedID(clientID, rqResource)
+		return token, err
+	case aadpodid.ServicePrincipal:
+		tenantid := azureID.Spec.TenantID
+		klog.Infof("matched identityType:%v tenantid:%s clientid:%s resource:%s", idType, tenantid, utils.RedactClientID(clientID), rqResource)
+		secret, err := mc.KubeClient.GetSecret(&azureID.Spec.ClientPassword)
+		if err != nil {
+			return nil, err
+		}
+		clientSecret := ""
+		for _, v := range secret.Data {
+			clientSecret = string(v)
+			break
+		}
+		token, err := auth.GetServicePrincipalToken(tenantid, clientID, clientSecret, rqResource)
+		return token, err
+	default:
+		return nil, fmt.Errorf("unsupported identity type %+v", idType)
+	}
+}
diff --git a/pkg/nmi/nmi.go b/pkg/nmi/nmi.go
new file mode 100644
index 00000000..13692a72
--- /dev/null
+++ b/pkg/nmi/nmi.go
@@ -0,0 +1,72 @@
+package nmi
+
+import (
+	"context"
+	"fmt"
+
+	aadpodid "github.com/Azure/aad-pod-identity/pkg/apis/aadpodidentity"
+	"github.com/Azure/aad-pod-identity/pkg/k8s"
+	"github.com/Azure/go-autorest/autorest/adal"
+	"k8s.io/klog"
+)
+
+// OperationMode is the mode in which NMI is operating
+// allowed values: standard, managed
+type OperationMode string
+
+// Config is the parameters used by token client
+type Config struct {
+	// Mode is the operation mode for token client
+	Mode string
+	// RetryAttemptsForCreated number of retries in NMI to find assigned identity in CREATED state for standard mode
+	RetryAttemptsForCreated int
+	// RetryAttemptsForAssigned number of retries in NMI to find assigned identity in ASSIGNED state for standard mode
+	RetryAttemptsForAssigned int
+	// FindIdentityRetryIntervalInSeconds Retry interval to find assigned identities in seconds for standard mode
+	FindIdentityRetryIntervalInSeconds int
+	// NodeName is the node on which NMI is running
+	NodeName string
+	// Namespaced makes NMI looks for identities in same namespace as pods
+	Namespaced bool
+}
+
+const (
+	// StandardMode ...
+	StandardMode OperationMode = "standard"
+	// ManagedMode ...
+	ManagedMode OperationMode = "managed"
+)
+
+// TokenClient ...
+type TokenClient interface {
+	// GetIdentities gets the list of identities which match the
+	// given pod in the form of AzureIdentity.
+	GetIdentities(ctx context.Context, podns, podname, clientID string) (*aadpodid.AzureIdentity, error)
+	// GetToken acquires a token by using the AzureIdentity.
+	GetToken(ctx context.Context, clientID, resource string, podID aadpodid.AzureIdentity) (token *adal.Token, err error)
+}
+
+// GetTokenClient returns a token client
+func GetTokenClient(client k8s.Client, config Config) (TokenClient, error) {
+	klog.Infof("Initializing in %s mode", config.Mode)
+
+	switch getOperationMode(config.Mode) {
+	case StandardMode:
+		return NewStandardTokenClient(client, config)
+	case ManagedMode:
+		return NewManagedTokenClient(client, config)
+	default:
+		return nil, fmt.Errorf("operation mode %s not supported", config.Mode)
+	}
+}
+
+func getOperationMode(mode string) OperationMode {
+	return OperationMode(mode)
+}
+
+// GetKubeClient returns kube client based on nmi mode
+func GetKubeClient(nodeName, mode string, enableScaleFeatures bool) (k8s.Client, error) {
+	// StandardMode client doesn't require azure identity and binding informers
+	// ManagedMode client doesn't require azure assigned identity informers
+	return k8s.NewKubeClient(nodeName, enableScaleFeatures, OperationMode(mode) == StandardMode)
+}
diff --git a/pkg/nmi/server/server.go b/pkg/nmi/server/server.go
index 90a48980..cb3ab7a9 100644
--- a/pkg/nmi/server/server.go
+++ b/pkg/nmi/server/server.go
@@ -1,7 +1,6 @@
 package server
 
 import (
-	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -18,13 +17,12 @@ import (
 	"syscall"
 	"time"
 
-	aadpodid "github.com/Azure/aad-pod-identity/pkg/apis/aadpodidentity"
 	auth "github.com/Azure/aad-pod-identity/pkg/auth"
 	k8s "github.com/Azure/aad-pod-identity/pkg/k8s"
 	"github.com/Azure/aad-pod-identity/pkg/metrics"
-	iptables "github.com/Azure/aad-pod-identity/pkg/nmi/iptables"
+	"github.com/Azure/aad-pod-identity/pkg/nmi"
+	"github.com/Azure/aad-pod-identity/pkg/nmi/iptables"
 	"github.com/Azure/aad-pod-identity/pkg/pod"
-	utils "github.com/Azure/aad-pod-identity/pkg/utils"
 	"github.com/Azure/go-autorest/autorest/adal"
 	"k8s.io/klog"
 )
@@ -43,15 +41,13 @@ type Server struct {
 	HostIP                             string
 	NodeName                           string
 	IPTableUpdateTimeIntervalInSeconds int
-	IsNamespaced                       bool
 	MICNamespace                       string
 	Initialized                        bool
 	BlockInstanceMetadata              bool
+	// TokenClient is client that fetches identities and tokens
+	TokenClient nmi.TokenClient
 
-	ListPodIDsRetryAttemptsForCreated  int
-	ListPodIDsRetryAttemptsForAssigned int
-	ListPodIDsRetryIntervalInSeconds   int
-	Reporter                           *metrics.Reporter
+	Reporter *metrics.Reporter
 }
 
 // NMIResponse is the response returned to caller
@@ -61,7 +57,7 @@ type NMIResponse struct {
 }
 
 // NewServer will create a new Server with default values.
-func NewServer(isNamespaced bool, micNamespace string, blockInstanceMetadata bool) *Server {
+func NewServer(micNamespace string, blockInstanceMetadata bool) *Server {
 	reporter, err := metrics.NewReporter()
 	if err != nil {
 		klog.Errorf("Error creating new reporter to emit metrics: %v", err)
@@ -71,7 +67,6 @@ func NewServer(isNamespaced bool, micNamespace string, blockInstanceMetadata boo
 		auth.InitReporter(reporter)
 	}
 	return &Server{
-		IsNamespaced:          isNamespaced,
 		MICNamespace:          micNamespace,
 		BlockInstanceMetadata: blockInstanceMetadata,
 		Reporter:              reporter,
@@ -186,7 +181,7 @@ func (fn appHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	rw := newResponseWriter(w)
 	ns := fn(rw, r)
 	latency := time.Since(start)
-	klog.Infof("Status (%d) took %d ns", rw.statusCode, latency.Nanoseconds())
+	klog.Infof("Status (%d) took %d ns for %s", rw.statusCode, latency.Nanoseconds(), tracker)
 
 	_, resource := parseRequestClientIDAndResource(r)
 
@@ -222,33 +217,13 @@ func (s *Server) hostHandler(w http.ResponseWriter, r *http.Request) (ns string)
 		http.Error(w, "parameter resource cannot be empty", http.StatusBadRequest)
 		return
 	}
-	podIDs, identityInCreatedStateFound, err := s.listPodIDsWithRetry(r.Context(), s.KubeClient, podns, podname, rqClientID)
+	podID, err := s.TokenClient.GetIdentities(r.Context(), podns, podname, rqClientID)
 	if err != nil {
-		msg := fmt.Sprintf("no AzureAssignedIdentity found for pod:%s/%s in desired state", podns, podname)
-		klog.Errorf("%s, %+v", msg, err)
-		http.Error(w, msg, getErrorResponseStatusCode(identityInCreatedStateFound))
+		klog.Error(err)
+		http.Error(w, err.Error(), getErrorResponseStatusCode(podID != nil))
 		return
 	}
-
-	// filter out if we are in namespaced mode
-	filterPodIdentities := []aadpodid.AzureIdentity{}
-	for _, val := range podIDs {
-		if s.IsNamespaced || aadpodid.IsNamespacedIdentity(&val) {
-			// namespaced mode
-			if val.Namespace == podns {
-				// matched namespace
-				filterPodIdentities = append(filterPodIdentities, val)
-			} else {
-				// unmatched namespaced
-				klog.Errorf("pod:%s/%s has identity %s/%s but identity is namespaced will be ignored", podns, podname, val.Name, val.Namespace)
-			}
-		} else {
-			// not in namespaced mode
-			filterPodIdentities = append(filterPodIdentities, val)
-		}
-	}
-	podIDs = filterPodIdentities
-	token, clientID, err := getTokenForMatchingID(s.KubeClient, rqClientID, rqResource, podIDs)
+	token, err := s.TokenClient.GetToken(r.Context(), rqClientID, rqResource, *podID)
 	if err != nil {
 		klog.Errorf("failed to get service principal token for pod:%s/%s, err: %+v", podns, podname, err)
 		http.Error(w, err.Error(), http.StatusForbidden)
@@ -256,7 +231,7 @@ func (s *Server) hostHandler(w http.ResponseWriter, r *http.Request) (ns string)
 	}
 	nmiResp := NMIResponse{
 		Token:    newMSIResponse(*token),
-		ClientID: clientID,
+		ClientID: podID.Spec.ClientID,
 	}
 	response, err := json.Marshal(nmiResp)
 	if err != nil {
@@ -374,15 +349,14 @@ func (s *Server) msiHandler(w http.ResponseWriter, r *http.Request) (ns string)
 		return
 	}
 
-	podIDs, identityInCreatedStateFound, err := s.listPodIDsWithRetry(r.Context(), s.KubeClient, podns, podname, rqClientID)
+	podID, err := s.TokenClient.GetIdentities(r.Context(), podns, podname, rqClientID)
 	if err != nil {
-		msg := fmt.Sprintf("no AzureAssignedIdentity found for pod:%s/%s in assigned state", podns, podname)
-		klog.Errorf("%s, %+v", msg, err)
-		http.Error(w, msg, getErrorResponseStatusCode(identityInCreatedStateFound))
+		klog.Error(err)
+		http.Error(w, err.Error(), getErrorResponseStatusCode(podID != nil))
 		return
 	}
 
-	token, _, err := getTokenForMatchingID(s.KubeClient, rqClientID, rqResource, podIDs)
+	token, err := s.TokenClient.GetToken(r.Context(), rqClientID, rqResource, *podID)
 	if err != nil {
 		klog.Errorf("failed to get service principal token for pod:%s/%s, %+v", podns, podname, err)
 		http.Error(w, err.Error(), http.StatusForbidden)
@@ -398,44 +372,6 @@ func (s *Server) msiHandler(w http.ResponseWriter, r *http.Request) (ns string)
 	return
 }
 
-func getTokenForMatchingID(kubeClient k8s.Client, rqClientID, rqResource string, podIDs []aadpodid.AzureIdentity) (token *adal.Token, clientID string, err error) {
-	rqHasClientID := len(rqClientID) != 0
-	for _, v := range podIDs {
-		clientID := v.Spec.ClientID
-		if rqHasClientID && !strings.EqualFold(rqClientID, clientID) {
-			klog.Warningf("clientid mismatch, requested:%s available:%s", rqClientID, clientID)
-			continue
-		}
-
-		idType := v.Spec.Type
-		switch idType {
-		case aadpodid.UserAssignedMSI:
-			klog.Infof("matched identityType:%v clientid:%s resource:%s", idType, utils.RedactClientID(clientID), rqResource)
-			token, err := auth.GetServicePrincipalTokenFromMSIWithUserAssignedID(clientID, rqResource)
-			return token, clientID, err
-		case aadpodid.ServicePrincipal:
-			tenantid := v.Spec.TenantID
-			klog.Infof("matched identityType:%v tenantid:%s clientid:%s resource:%s", idType, tenantid, utils.RedactClientID(clientID), rqResource)
-			secret, err := kubeClient.GetSecret(&v.Spec.ClientPassword)
-			if err != nil {
-				return nil, clientID, err
-			}
-			clientSecret := ""
-			for _, v := range secret.Data {
-				clientSecret = string(v)
-				break
-			}
-			token, err := auth.GetServicePrincipalToken(tenantid, clientID, clientSecret, rqResource)
-			return token, clientID, err
-		default:
-			return nil, clientID, fmt.Errorf("unsupported identity type %+v", idType)
-		}
-	}
-
-	// We have not yet returned, so pass up an error
-	return nil, "", fmt.Errorf("azureidentity is not configured for the pod")
-}
-
 func parseRequestHeader(r *http.Request) (podns string, podname string) {
 	podns = r.Header.Get("podns")
 	podname = r.Header.Get("podname")
@@ -531,75 +467,6 @@ func handleTermination() {
 	os.Exit(exitCode)
 }
 
-// listPodIDsWithRetry returns a list of matched identities in Assigned state, boolean indicating if at least an identity was found in Created state and error if any
-func (s *Server) listPodIDsWithRetry(ctx context.Context, kubeClient k8s.Client, podns, podname, rqClientID string) ([]aadpodid.AzureIdentity, bool, error) {
-	attempt := 0
-	var err error
-	var idStateMap map[string][]aadpodid.AzureIdentity
-
-	// this loop will run to ensure we have assigned identities before we return. If there are no assigned identities in created state within 80s (16 retries * 5s wait) then we return an error.
-	// If we get an assigned identity in created state within 80s, then loop will continue until 100s to find assigned identity in assigned state.
-	// Retry interval for CREATED state is set to 80s because avg time for identity to be assigned to the node is 35-37s.
-	for attempt < s.ListPodIDsRetryAttemptsForCreated+s.ListPodIDsRetryAttemptsForAssigned {
-		idStateMap, err = kubeClient.ListPodIds(podns, podname)
-		if err == nil {
-			if len(rqClientID) == 0 {
-				// check to ensure backward compatability with assignedIDs that have no state
-				// assigned identites created with old version of mic will not contain a state. So first we check to see if an assigned identity with
-				// no state exists that matches req client id.
-				if len(idStateMap[""]) != 0 {
-					klog.Warningf("found assignedIDs with no state for pod:%s/%s. AssignedIDs created with old version of mic.", podns, podname)
-					return idStateMap[""], true, nil
-				}
-				if len(idStateMap[aadpodid.AssignedIDAssigned]) != 0 {
-					return idStateMap[aadpodid.AssignedIDAssigned], true, nil
-				}
-				if len(idStateMap[aadpodid.AssignedIDCreated]) == 0 && attempt >= s.ListPodIDsRetryAttemptsForCreated {
-					return nil, false, fmt.Errorf("getting assigned identities for pod %s/%s in CREATED state failed after %d attempts, retry duration [%d]s. Error: %v",
-						podns, podname, s.ListPodIDsRetryAttemptsForCreated, s.ListPodIDsRetryIntervalInSeconds, err)
-				}
-			} else {
-				// if client id exists in request, we need to ensure the identity with this client
-				// exists and is in Assigned state
-				// check to ensure backward compatability with assignedIDs that have no state
-				for _, podID := range idStateMap[""] {
-					if strings.EqualFold(rqClientID, podID.Spec.ClientID) {
-						klog.Warningf("found assignedIDs with no state for pod:%s/%s. AssignedIDs created with old version of mic.", podns, podname)
-						return idStateMap[""], true, nil
-					}
-				}
-				for _, podID := range idStateMap[aadpodid.AssignedIDAssigned] {
-					if strings.EqualFold(rqClientID, podID.Spec.ClientID) {
-						return idStateMap[aadpodid.AssignedIDAssigned], true, nil
-					}
-				}
-				var foundMatch bool
-				for _, podID := range idStateMap[aadpodid.AssignedIDCreated] {
-					if strings.EqualFold(rqClientID, podID.Spec.ClientID) {
-						foundMatch = true
-						break
-					}
-				}
-				if !foundMatch && attempt >= s.ListPodIDsRetryAttemptsForCreated {
-					return nil, false, fmt.Errorf("getting assigned identities for pod %s/%s in CREATED state failed after %d attempts, retry duration [%d]s. Error: %v",
-						podns, podname, s.ListPodIDsRetryAttemptsForCreated, s.ListPodIDsRetryIntervalInSeconds, err)
-				}
-			}
-		}
-		attempt++
-
-		select {
-		case <-time.After(time.Duration(s.ListPodIDsRetryIntervalInSeconds) * time.Second):
-		case <-ctx.Done():
-			err = ctx.Err()
-			return nil, true, err
-		}
-		klog.V(4).Infof("failed to get assigned ids for pod:%s/%s in ASSIGNED state, retrying attempt: %d", podns, podname, attempt)
-	}
-	return nil, true, fmt.Errorf("getting assigned identities for pod %s/%s in ASSIGNED state failed after %d attempts, retry duration [%d]s. Error: %v",
-		podns, podname, s.ListPodIDsRetryAttemptsForCreated+s.ListPodIDsRetryAttemptsForAssigned, s.ListPodIDsRetryIntervalInSeconds, err)
-}
-
 func getErrorResponseStatusCode(identityFound bool) int {
 	// if at least an identity was found in created state then we return 404 which is a retriable error code
 	// in the go-autorest library. If the identity is in CREATED state then the identity is being processed in
diff --git a/pkg/nmi/standard.go b/pkg/nmi/standard.go
new file mode 100644
index 00000000..27a74f3c
--- /dev/null
+++ b/pkg/nmi/standard.go
@@ -0,0 +1,184 @@
+package nmi
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"time"
+
+	aadpodid "github.com/Azure/aad-pod-identity/pkg/apis/aadpodidentity"
+	auth "github.com/Azure/aad-pod-identity/pkg/auth"
+	k8s "github.com/Azure/aad-pod-identity/pkg/k8s"
+	utils "github.com/Azure/aad-pod-identity/pkg/utils"
+	"github.com/Azure/go-autorest/autorest/adal"
+
+	"k8s.io/klog"
+)
+
+// StandardClient implements the TokenClient interface
+type StandardClient struct {
+	TokenClient
+	KubeClient                         k8s.Client
+	ListPodIDsRetryAttemptsForCreated  int
+	ListPodIDsRetryAttemptsForAssigned int
+	ListPodIDsRetryIntervalInSeconds   int
+	IsNamespaced                       bool
+}
+
+// NewStandardTokenClient creates new standard nmi client
+func NewStandardTokenClient(client k8s.Client, config Config) (*StandardClient, error) {
+	return &StandardClient{
+		KubeClient:                         client,
+		ListPodIDsRetryAttemptsForCreated:  config.RetryAttemptsForCreated,
+		ListPodIDsRetryAttemptsForAssigned: config.RetryAttemptsForAssigned,
+		ListPodIDsRetryIntervalInSeconds:   config.FindIdentityRetryIntervalInSeconds,
+		IsNamespaced:                       config.Namespaced,
+	}, nil
+}
+
+// GetIdentities ...
+func (sc *StandardClient) GetIdentities(ctx context.Context, podns, podname, clientID string) (*aadpodid.AzureIdentity, error) {
+	podIDs, identityInCreatedStateFound, err := sc.listPodIDsWithRetry(ctx, podns, podname, clientID)
+	if err != nil {
+		// if identity not found in created state return nil identity which is then used to send 403 error
+		if !identityInCreatedStateFound {
+			return nil, err
+		}
+		// identity found in created state but there was an error, then return empty struct which will result in 404 error
+		return &aadpodid.AzureIdentity{}, err
+	}
+
+	// filter out if we are in namespaced mode
+	filterPodIdentities := []aadpodid.AzureIdentity{}
+	for _, val := range podIDs {
+		if sc.IsNamespaced || aadpodid.IsNamespacedIdentity(&val) {
+			// namespaced mode
+			if val.Namespace == podns {
+				// matched namespace
+				filterPodIdentities = append(filterPodIdentities, val)
+			} else {
+				// unmatched namespaced
+				klog.Errorf("pod:%s/%s has identity %s/%s but identity is namespaced will be ignored", podns, podname, val.Name, val.Namespace)
+			}
+		} else {
+			// not in namespaced mode
+			filterPodIdentities = append(filterPodIdentities, val)
+		}
+	}
+
+	for _, id := range filterPodIdentities {
+		// if client doesn't exist in the request, then return the first identity
+		if len(clientID) == 0 {
+			klog.Infof("No clientID in request. %s/%s has been matched with azure identity %s/%s", podns, podname, id.Namespace, id.Name)
+			return &id, nil
+		}
+		// if client id exists in the request, then send the first identity that matched the client id
+		if len(clientID) != 0 && id.Spec.ClientID == clientID {
+			return &id, nil
+		}
+	}
+
+	return nil, fmt.Errorf("azureidentity is not configured for the pod")
+}
+
+// listPodIDsWithRetry returns a list of matched identities in Assigned state, boolean indicating if at least an identity was found in Created state and error if any
+func (sc *StandardClient) listPodIDsWithRetry(ctx context.Context, podns, podname, rqClientID string) ([]aadpodid.AzureIdentity, bool, error) {
+	attempt := 0
+	var err error
+	var idStateMap map[string][]aadpodid.AzureIdentity
+
+	// this loop will run to ensure we have assigned identities before we return. If there are no assigned identities in created state within 80s (16 retries * 5s wait) then we return an error.
+	// If we get an assigned identity in created state within 80s, then loop will continue until 100s to find assigned identity in assigned state.
+	// Retry interval for CREATED state is set to 80s because avg time for identity to be assigned to the node is 35-37s.
+	for attempt < sc.ListPodIDsRetryAttemptsForCreated+sc.ListPodIDsRetryAttemptsForAssigned {
+		idStateMap, err = sc.KubeClient.ListPodIds(podns, podname)
+		if err == nil {
+			if len(rqClientID) == 0 {
+				// check to ensure backward compatability with assignedIDs that have no state
+				// assigned identites created with old version of mic will not contain a state. So first we check to see if an assigned identity with
+				// no state exists that matches req client id.
+				if len(idStateMap[""]) != 0 {
+					klog.Warningf("found assignedIDs with no state for pod:%s/%s. AssignedIDs created with old version of mic.", podns, podname)
+					return idStateMap[""], true, nil
+				}
+				if len(idStateMap[aadpodid.AssignedIDAssigned]) != 0 {
+					return idStateMap[aadpodid.AssignedIDAssigned], true, nil
+				}
+				if len(idStateMap[aadpodid.AssignedIDCreated]) == 0 && attempt >= sc.ListPodIDsRetryAttemptsForCreated {
+					return nil, false, fmt.Errorf("getting assigned identities for pod %s/%s in CREATED state failed after %d attempts, retry duration [%d]s. Error: %v",
+						podns, podname, sc.ListPodIDsRetryAttemptsForCreated, sc.ListPodIDsRetryIntervalInSeconds, err)
+				}
+			} else {
+				// if client id exists in request, we need to ensure the identity with this client
+				// exists and is in Assigned state
+				// check to ensure backward compatability with assignedIDs that have no state
+				for _, podID := range idStateMap[""] {
+					if strings.EqualFold(rqClientID, podID.Spec.ClientID) {
+						klog.Warningf("found assignedIDs with no state for pod:%s/%s. AssignedIDs created with old version of mic.", podns, podname)
+						return idStateMap[""], true, nil
+					}
+				}
+				for _, podID := range idStateMap[aadpodid.AssignedIDAssigned] {
+					if strings.EqualFold(rqClientID, podID.Spec.ClientID) {
+						return idStateMap[aadpodid.AssignedIDAssigned], true, nil
+					}
+				}
+				var foundMatch bool
+				for _, podID := range idStateMap[aadpodid.AssignedIDCreated] {
+					if strings.EqualFold(rqClientID, podID.Spec.ClientID) {
+						foundMatch = true
+						break
+					}
+				}
+				if !foundMatch && attempt >= sc.ListPodIDsRetryAttemptsForCreated {
+					return nil, false, fmt.Errorf("getting assigned identities for pod %s/%s in CREATED state failed after %d attempts, retry duration [%d]s. Error: %v",
+						podns, podname, sc.ListPodIDsRetryAttemptsForCreated, sc.ListPodIDsRetryIntervalInSeconds, err)
+				}
+			}
+		}
+		attempt++
+
+		select {
+		case <-time.After(time.Duration(sc.ListPodIDsRetryIntervalInSeconds) * time.Second):
+		case <-ctx.Done():
+			err = ctx.Err()
+			return nil, true, err
+		}
+		klog.V(4).Infof("failed to get assigned ids for pod:%s/%s in ASSIGNED state, retrying attempt: %d", podns, podname, attempt)
+	}
+	return nil, true, fmt.Errorf("getting assigned identities for pod %s/%s in ASSIGNED state failed after %d attempts, retry duration [%d]s. Error: %v",
+		podns, podname, sc.ListPodIDsRetryAttemptsForCreated+sc.ListPodIDsRetryAttemptsForAssigned, sc.ListPodIDsRetryIntervalInSeconds, err)
+}
+
+// GetToken ...
+func (sc *StandardClient) GetToken(ctx context.Context, rqClientID, rqResource string, azureID aadpodid.AzureIdentity) (token *adal.Token, err error) {
+	rqHasClientID := len(rqClientID) != 0
+	clientID := azureID.Spec.ClientID
+
+	idType := azureID.Spec.Type
+	switch idType {
+	case aadpodid.UserAssignedMSI:
+		if rqHasClientID && !strings.EqualFold(rqClientID, clientID) {
+			klog.Warningf("clientid mismatch, requested:%s available:%s", rqClientID, clientID)
+		}
+		klog.Infof("matched identityType:%v clientid:%s resource:%s", idType, utils.RedactClientID(clientID), rqResource)
+		token, err := auth.GetServicePrincipalTokenFromMSIWithUserAssignedID(clientID, rqResource)
+		return token, err
+	case aadpodid.ServicePrincipal:
+		tenantid := azureID.Spec.TenantID
+		klog.Infof("matched identityType:%v tenantid:%s clientid:%s resource:%s", idType, tenantid, utils.RedactClientID(clientID), rqResource)
+		secret, err := sc.KubeClient.GetSecret(&azureID.Spec.ClientPassword)
+		if err != nil {
+			return nil, err
+		}
+		clientSecret := ""
+		for _, v := range secret.Data {
+			clientSecret = string(v)
+			break
+		}
+		token, err := auth.GetServicePrincipalToken(tenantid, clientID, clientSecret, rqResource)
+		return token, err
+	default:
+		return nil, fmt.Errorf("unsupported identity type %+v", idType)
+	}
+}


Test patch:
diff --git a/pkg/mic/mic_test.go b/pkg/mic/mic_test.go
index 176f2b92..fc02361c 100644
--- a/pkg/mic/mic_test.go
+++ b/pkg/mic/mic_test.go
@@ -11,6 +11,7 @@ import (
 	internalaadpodid "github.com/Azure/aad-pod-identity/pkg/apis/aadpodidentity"
 	aadpodid "github.com/Azure/aad-pod-identity/pkg/apis/aadpodidentity/v1"
 	"github.com/Azure/aad-pod-identity/pkg/config"
+	"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"
@@ -21,6 +22,7 @@ import (
 	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/client-go/rest"
+	"k8s.io/client-go/tools/cache"
 	"k8s.io/klog"
 )
 
@@ -327,7 +329,7 @@ func (c *TestPodClient) DeletePod(podName string, podNs string) {
 /****************** CRD MOCK ****************************/
 
 type TestCrdClient struct {
-	*Client
+	*crd.Client
 	mu            sync.Mutex
 	assignedIDMap map[string]*internalaadpodid.AzureAssignedIdentity
 	bindingMap    map[string]*aadpodid.AzureIdentityBinding
@@ -346,11 +348,7 @@ func NewTestCrdClient(config *rest.Config) *TestCrdClient {
 func (c *TestCrdClient) Start(exit <-chan struct{}) {
 }
 
-func (c *TestCrdClient) SyncCache(exit <-chan struct{}) {
-
-}
-
-func (c *TestCrdClient) SyncCacheLite(exit <-chan struct{}) {
+func (c *TestCrdClient) SyncCache(exit <-chan struct{}, initial bool, cacheSyncs ...cache.InformerSynced) {
 
 }
 
diff --git a/pkg/nmi/managed_test.go b/pkg/nmi/managed_test.go
new file mode 100644
index 00000000..d3cf121a
--- /dev/null
+++ b/pkg/nmi/managed_test.go
@@ -0,0 +1,139 @@
+package nmi
+
+import (
+	"context"
+	"reflect"
+	"testing"
+
+	aadpodid "github.com/Azure/aad-pod-identity/pkg/apis/aadpodidentity"
+	v1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+)
+
+func (c *TestKubeClient) ListPodIdsWithBinding(podns string, labels map[string]string) ([]aadpodid.AzureIdentity, error) {
+	identities, _ := c.azureIdentities.([]aadpodid.AzureIdentity)
+	return identities, nil
+}
+
+func (c *TestKubeClient) GetPod(ns, name string) (v1.Pod, error) {
+	return v1.Pod{}, nil
+}
+
+func TestGetIdentitiesManagedClient(t *testing.T) {
+	cases := []struct {
+		name                  string
+		azureIdentities       []aadpodid.AzureIdentity
+		clientID              string
+		expectedErr           bool
+		expectedAzureIdentity *aadpodid.AzureIdentity
+		isNamespaced          bool
+		podName               string
+		podNamespace          string
+	}{
+		{
+			name:                  "no azure identity found",
+			azureIdentities:       nil,
+			expectedErr:           true,
+			expectedAzureIdentity: nil,
+			podName:               "pod1",
+			podNamespace:          "default",
+		},
+		{
+			name: "clientID in request, but no matching identity",
+			azureIdentities: []aadpodid.AzureIdentity{
+				{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "azid2",
+						Namespace: "default",
+					},
+					Spec: aadpodid.AzureIdentitySpec{
+						ClientID: "clientid2",
+					},
+				},
+			},
+			expectedErr:           true,
+			expectedAzureIdentity: nil,
+			podName:               "pod2",
+			podNamespace:          "default",
+			clientID:              "clientid1",
+		},
+		{
+			name: "clientID in request, found matching identity",
+			azureIdentities: []aadpodid.AzureIdentity{
+				{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "azid3",
+						Namespace: "default",
+					},
+					Spec: aadpodid.AzureIdentitySpec{
+						ClientID: "clientid3",
+					},
+				},
+			},
+			expectedErr: false,
+			expectedAzureIdentity: &aadpodid.AzureIdentity{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "azid3",
+					Namespace: "default",
+				},
+				Spec: aadpodid.AzureIdentitySpec{
+					ClientID: "clientid3",
+				},
+			},
+			podName:      "pod3",
+			podNamespace: "default",
+			clientID:     "clientid3",
+		},
+		{
+			name: "no clientID in request, first matching identity in namespace returned",
+			azureIdentities: []aadpodid.AzureIdentity{
+				{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "azid2",
+						Namespace: "default",
+					},
+					Spec: aadpodid.AzureIdentitySpec{
+						ClientID: "clientid2",
+					},
+				},
+				{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "azid3",
+						Namespace: "default",
+					},
+					Spec: aadpodid.AzureIdentitySpec{
+						ClientID: "clientid3",
+					},
+				},
+			},
+			expectedErr: false,
+			expectedAzureIdentity: &aadpodid.AzureIdentity{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "azid2",
+					Namespace: "default",
+				},
+				Spec: aadpodid.AzureIdentitySpec{
+					ClientID: "clientid2",
+				},
+			},
+			podName:      "pod4",
+			podNamespace: "default",
+		},
+	}
+
+	for i, tc := range cases {
+		t.Log(i, tc.name)
+		tokenClient, err := NewManagedTokenClient(NewTestKubeClient(tc.azureIdentities), Config{Namespaced: true})
+		if err != nil {
+			t.Fatalf("expected err to be nil, got: %v", err)
+		}
+
+		azIdentity, err := tokenClient.GetIdentities(context.Background(), tc.podNamespace, tc.podName, tc.clientID)
+		if tc.expectedErr != (err != nil) {
+			t.Fatalf("expected error: %v, got: %v", tc.expectedErr, err)
+		}
+		if !reflect.DeepEqual(tc.expectedAzureIdentity, azIdentity) {
+			t.Fatalf("expected the azure identity to be equal")
+		}
+	}
+}
diff --git a/pkg/nmi/server/server_test.go b/pkg/nmi/server/server_test.go
deleted file mode 100644
index 97756843..00000000
--- a/pkg/nmi/server/server_test.go
+++ /dev/null
@@ -1,47 +0,0 @@
-package server
-
-import (
-	"encoding/base64"
-	"testing"
-
-	internalaadpodid "github.com/Azure/aad-pod-identity/pkg/apis/aadpodidentity"
-	auth "github.com/Azure/aad-pod-identity/pkg/auth"
-	"github.com/Azure/aad-pod-identity/pkg/k8s"
-	"github.com/Azure/aad-pod-identity/pkg/metrics"
-
-	v1 "k8s.io/api/core/v1"
-	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
-	"k8s.io/client-go/kubernetes/fake"
-)
-
-func TestGetTokenForMatchingIDBySP(t *testing.T) {
-	fakeClient := fake.NewSimpleClientset()
-	reporter, err := metrics.NewReporter()
-	if err != nil {
-		t.Fatalf("expected nil error, got: %+v", err)
-	}
-	auth.InitReporter(reporter)
-
-	secret := &v1.Secret{ObjectMeta: metav1.ObjectMeta{Name: "clientSecret"}, Data: make(map[string][]byte)}
-	val, _ := base64.StdEncoding.DecodeString("YWJjZA==")
-	secret.Data["key1"] = val
-	fakeClient.CoreV1().Secrets("default").Create(secret)
-
-	kubeClient := &k8s.KubeClient{ClientSet: fakeClient}
-
-	secretRef := v1.SecretReference{
-		Name:      "clientSecret",
-		Namespace: "default",
-	}
-
-	podID := internalaadpodid.AzureIdentity{
-		Spec: internalaadpodid.AzureIdentitySpec{
-			Type:           internalaadpodid.ServicePrincipal,
-			TenantID:       "tid",
-			ClientID:       "aabc0000-a83v-9h4m-000j-2c0a66b0c1f9",
-			ClientPassword: secretRef,
-		},
-	}
-	podIDs := []internalaadpodid.AzureIdentity{podID}
-	getTokenForMatchingID(kubeClient, podID.Spec.ClientID, "https://management.azure.com/", podIDs)
-}
diff --git a/pkg/nmi/standard_test.go b/pkg/nmi/standard_test.go
new file mode 100644
index 00000000..3d427949
--- /dev/null
+++ b/pkg/nmi/standard_test.go
@@ -0,0 +1,347 @@
+package nmi
+
+import (
+	"context"
+	"encoding/base64"
+	"reflect"
+	"testing"
+
+	aadpodid "github.com/Azure/aad-pod-identity/pkg/apis/aadpodidentity"
+	auth "github.com/Azure/aad-pod-identity/pkg/auth"
+	"github.com/Azure/aad-pod-identity/pkg/k8s"
+	"github.com/Azure/aad-pod-identity/pkg/metrics"
+
+	v1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/client-go/kubernetes/fake"
+)
+
+type TestKubeClient struct {
+	k8s.Client
+	azureIdentities interface{}
+	err             error
+}
+
+func NewTestKubeClient(azids interface{}) *TestKubeClient {
+	return &TestKubeClient{
+		azureIdentities: azids,
+	}
+}
+
+func (c *TestKubeClient) setError(err error) {
+	c.err = err
+}
+
+func (c *TestKubeClient) ListPodIds(podns, podname string) (map[string][]aadpodid.AzureIdentity, error) {
+	identities, _ := c.azureIdentities.(map[string][]aadpodid.AzureIdentity)
+	return identities, c.err
+}
+
+func TestGetTokenForMatchingIDBySP(t *testing.T) {
+	fakeClient := fake.NewSimpleClientset()
+	reporter, err := metrics.NewReporter()
+	if err != nil {
+		t.Fatalf("expected nil error, got: %+v", err)
+	}
+	auth.InitReporter(reporter)
+
+	secret := &v1.Secret{ObjectMeta: metav1.ObjectMeta{Name: "clientSecret"}, Data: make(map[string][]byte)}
+	val, _ := base64.StdEncoding.DecodeString("YWJjZA==")
+	secret.Data["key1"] = val
+	fakeClient.CoreV1().Secrets("default").Create(secret)
+
+	kubeClient := &k8s.KubeClient{ClientSet: fakeClient}
+	tokenClient, err := NewStandardTokenClient(kubeClient, Config{})
+	if err != nil {
+		t.Fatalf("expected err to be nil, got: %v", err)
+	}
+
+	secretRef := v1.SecretReference{
+		Name:      "clientSecret",
+		Namespace: "default",
+	}
+
+	podID := aadpodid.AzureIdentity{
+		Spec: aadpodid.AzureIdentitySpec{
+			Type:           aadpodid.ServicePrincipal,
+			TenantID:       "tid",
+			ClientID:       "aabc0000-a83v-9h4m-000j-2c0a66b0c1f9",
+			ClientPassword: secretRef,
+		},
+	}
+	tokenClient.GetToken(context.Background(), podID.Spec.ClientID, "https://management.azure.com/", podID)
+}
+
+func TestGetIdentitiesStandardClient(t *testing.T) {
+	cases := []struct {
+		name                  string
+		azureIdentities       map[string][]aadpodid.AzureIdentity
+		clientID              string
+		expectedErr           bool
+		expectedAzureIdentity *aadpodid.AzureIdentity
+		isNamespaced          bool
+		podName               string
+		podNamespace          string
+	}{
+		{
+			name:                  "no azure identities",
+			azureIdentities:       make(map[string][]aadpodid.AzureIdentity),
+			expectedErr:           true,
+			expectedAzureIdentity: nil,
+			podName:               "pod1",
+			podNamespace:          "default",
+		},
+		{
+			name: "azure identities with old 1.3/1.4, no request client id",
+			azureIdentities: map[string][]aadpodid.AzureIdentity{
+				"": []aadpodid.AzureIdentity{
+					{
+						ObjectMeta: metav1.ObjectMeta{
+							Name:      "azid1",
+							Namespace: "default",
+						},
+						Spec: aadpodid.AzureIdentitySpec{
+							ClientID: "clientid1",
+						},
+					},
+					{
+						ObjectMeta: metav1.ObjectMeta{
+							Name:      "azid2",
+							Namespace: "default",
+						},
+						Spec: aadpodid.AzureIdentitySpec{
+							ClientID: "clientid2",
+						},
+					},
+				},
+			},
+			expectedErr: false,
+			expectedAzureIdentity: &aadpodid.AzureIdentity{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "azid1",
+					Namespace: "default",
+				},
+				Spec: aadpodid.AzureIdentitySpec{
+					ClientID: "clientid1",
+				},
+			},
+			podName:      "pod2",
+			podNamespace: "default",
+		},
+		{
+			name: "no request client id, found in created state only",
+			azureIdentities: map[string][]aadpodid.AzureIdentity{
+				aadpodid.AssignedIDCreated: []aadpodid.AzureIdentity{
+					{
+						ObjectMeta: metav1.ObjectMeta{
+							Name:      "azid3",
+							Namespace: "default",
+						},
+						Spec: aadpodid.AzureIdentitySpec{
+							ClientID: "clientid3",
+						},
+					},
+					{
+						ObjectMeta: metav1.ObjectMeta{
+							Name:      "azid4",
+							Namespace: "default",
+						},
+						Spec: aadpodid.AzureIdentitySpec{
+							ClientID: "clientid4",
+						},
+					},
+				},
+			},
+			expectedAzureIdentity: &aadpodid.AzureIdentity{},
+			expectedErr:           true,
+			podName:               "pod3",
+			podNamespace:          "default",
+		},
+		{
+			name: "no request client id, found in assigned state",
+			azureIdentities: map[string][]aadpodid.AzureIdentity{
+				aadpodid.AssignedIDAssigned: []aadpodid.AzureIdentity{
+					{
+						ObjectMeta: metav1.ObjectMeta{
+							Name:      "azid5",
+							Namespace: "default",
+						},
+						Spec: aadpodid.AzureIdentitySpec{
+							ClientID: "clientid5",
+						},
+					},
+					{
+						ObjectMeta: metav1.ObjectMeta{
+							Name:      "azid6",
+							Namespace: "default",
+						},
+						Spec: aadpodid.AzureIdentitySpec{
+							ClientID: "clientid6",
+						},
+					},
+				},
+			},
+			expectedErr: false,
+			expectedAzureIdentity: &aadpodid.AzureIdentity{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "azid5",
+					Namespace: "default",
+				},
+				Spec: aadpodid.AzureIdentitySpec{
+					ClientID: "clientid5",
+				},
+			},
+			podName:      "pod4",
+			podNamespace: "default",
+		},
+		{
+			name: "client id in request, no identity with same client id in assigned state",
+			azureIdentities: map[string][]aadpodid.AzureIdentity{
+				aadpodid.AssignedIDCreated: []aadpodid.AzureIdentity{
+					{
+						ObjectMeta: metav1.ObjectMeta{
+							Name:      "azid1",
+							Namespace: "default",
+						},
+						Spec: aadpodid.AzureIdentitySpec{
+							ClientID: "clientid1",
+						},
+					},
+				},
+				aadpodid.AssignedIDAssigned: []aadpodid.AzureIdentity{
+					{
+						ObjectMeta: metav1.ObjectMeta{
+							Name:      "azid2",
+							Namespace: "default",
+						},
+						Spec: aadpodid.AzureIdentitySpec{
+							ClientID: "clientid2",
+						},
+					},
+				},
+			},
+			expectedErr:           true,
+			expectedAzureIdentity: &aadpodid.AzureIdentity{},
+			podName:               "pod5",
+			podNamespace:          "default",
+			clientID:              "clientid1",
+		},
+		{
+			name: "client id in request, identity in same namespace returned with force namespace mode",
+			azureIdentities: map[string][]aadpodid.AzureIdentity{
+				aadpodid.AssignedIDAssigned: []aadpodid.AzureIdentity{
+					{
+						ObjectMeta: metav1.ObjectMeta{
+							Name:      "azid2",
+							Namespace: "testns",
+						},
+						Spec: aadpodid.AzureIdentitySpec{
+							ClientID: "clientid2",
+						},
+					},
+					{
+						ObjectMeta: metav1.ObjectMeta{
+							Name:      "azid1",
+							Namespace: "default",
+						},
+						Spec: aadpodid.AzureIdentitySpec{
+							ClientID: "clientid1",
+						},
+					},
+					{
+						ObjectMeta: metav1.ObjectMeta{
+							Name:      "azid3",
+							Namespace: "default",
+						},
+						Spec: aadpodid.AzureIdentitySpec{
+							ClientID: "clientid3",
+						},
+					},
+				},
+			},
+			expectedErr: false,
+			expectedAzureIdentity: &aadpodid.AzureIdentity{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "azid2",
+					Namespace: "testns",
+				},
+				Spec: aadpodid.AzureIdentitySpec{
+					ClientID: "clientid2",
+				},
+			},
+			podName:      "pod7",
+			podNamespace: "testns",
+			clientID:     "clientid2",
+			isNamespaced: true,
+		},
+		{
+			name: "no client id in request, identity in same namespace returned with force namespace mode",
+			azureIdentities: map[string][]aadpodid.AzureIdentity{
+				aadpodid.AssignedIDAssigned: []aadpodid.AzureIdentity{
+					{
+						ObjectMeta: metav1.ObjectMeta{
+							Name:      "azid2",
+							Namespace: "default",
+						},
+						Spec: aadpodid.AzureIdentitySpec{
+							ClientID: "clientid2",
+						},
+					},
+					{
+						ObjectMeta: metav1.ObjectMeta{
+							Name:      "azid1",
+							Namespace: "default",
+						},
+						Spec: aadpodid.AzureIdentitySpec{
+							ClientID: "clientid1",
+						},
+					},
+					{
+						ObjectMeta: metav1.ObjectMeta{
+							Name:      "azid3",
+							Namespace: "testns",
+						},
+						Spec: aadpodid.AzureIdentitySpec{
+							ClientID: "clientid3",
+						},
+					},
+				},
+			},
+			expectedErr: false,
+			expectedAzureIdentity: &aadpodid.AzureIdentity{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "azid3",
+					Namespace: "testns",
+				},
+				Spec: aadpodid.AzureIdentitySpec{
+					ClientID: "clientid3",
+				},
+			},
+			podName:      "pod8",
+			podNamespace: "testns",
+			isNamespaced: true,
+		},
+	}
+
+	for i, tc := range cases {
+		t.Log(i, tc.name)
+		tokenClient, err := NewStandardTokenClient(NewTestKubeClient(tc.azureIdentities), Config{
+			Mode:                               "standard",
+			RetryAttemptsForCreated:            2,
+			RetryAttemptsForAssigned:           1,
+			FindIdentityRetryIntervalInSeconds: 1,
+			Namespaced:                         tc.isNamespaced,
+		})
+		if err != nil {
+			t.Fatalf("expected err to be nil, got: %v", err)
+		}
+
+		azIdentity, err := tokenClient.GetIdentities(context.Background(), tc.podNamespace, tc.podName, tc.clientID)
+		if tc.expectedErr != (err != nil) {
+			t.Fatalf("expected error: %v, got: %v", tc.expectedErr, err)
+		}
+		if !reflect.DeepEqual(tc.expectedAzureIdentity, azIdentity) {
+			t.Fatalf("expected the azure identity to be equal")
+		}
+	}
+}
diff --git a/test/common/k8s/azureassignedidentity/azureassignedidentity.go b/test/common/k8s/azureassignedidentity/azureassignedidentity.go
index 49e98a93..204bf21e 100644
--- a/test/common/k8s/azureassignedidentity/azureassignedidentity.go
+++ b/test/common/k8s/azureassignedidentity/azureassignedidentity.go
@@ -68,14 +68,14 @@ func GetAllByPrefix(prefix string) ([]aadpodid.AzureAssignedIdentity, error) {
 }
 
 // WaitOnLengthMatched will block until the number of Azure Assigned Identity matches the target
-func WaitOnLengthMatched(target int) (bool, error) {
+func WaitOnLengthMatched(target int, checkStatus bool) (bool, error) {
 	successChannel, errorChannel := make(chan bool, 1), make(chan error)
 	// defining ~2 mins 30 seconds as an acceptable timeframe for ids to be assigned to node
 	duration, sleep := 150*time.Second, 10*time.Second
 	ctx, cancel := context.WithTimeout(context.Background(), duration)
 	defer cancel()
 
-	fmt.Println("# Tight-poll to check if the Azure Assigned Identity is deleted...")
+	fmt.Printf("# Tight-poll to check if the # of Azure Assigned Identity is %d...\n", target)
 	go func() {
 		for {
 			select {
@@ -88,10 +88,28 @@ func WaitOnLengthMatched(target int) (bool, error) {
 					return
 				}
 				// len of nil slices is 0, so shouldn't panic here
-				if len(list.Items) == target {
+				if (!checkStatus || target == 0) && len(list.Items) == target {
 					successChannel <- true
 					return
 				}
+				// if checking for assigned identity count to be 0, then we just check if len matches
+				// but if checking for assigned identity count > 1, apart from checking if the len of
+				// assigned identities is desired, we also check if all the assigned identities are in
+				// desired state
+				if (checkStatus || target != 0) && list.Items != nil && len(list.Items) == target {
+					var totalAssigned int
+					for _, item := range list.Items {
+						if item.Status.Status != aadpodid.AssignedIDAssigned {
+							break
+						}
+						totalAssigned++
+					}
+					if totalAssigned == target {
+						successChannel <- true
+						return
+					}
+				}
+
 				fmt.Printf("# The number of Azure Assigned Identity is not equal to %d yet. Retrying in %s...\n", target, sleep.String())
 				time.Sleep(sleep)
 			}
diff --git a/test/common/k8s/infra/infra.go b/test/common/k8s/infra/infra.go
index fbf4a2e3..47ab9bfb 100644
--- a/test/common/k8s/infra/infra.go
+++ b/test/common/k8s/infra/infra.go
@@ -11,7 +11,7 @@ import (
 )
 
 // CreateInfra will deploy all the infrastructure components (nmi and mic) on a Kubernetes cluster
-func CreateInfra(namespace, registry, nmiVersion, micVersion, templateOutputPath string, old, enableScaleFeatures bool, immutableUserMSIs string) error {
+func CreateInfra(namespace, registry, nmiVersion, micVersion, templateOutputPath string, old, enableScaleFeatures bool, immutableUserMSIs, nmiMode string) error {
 	var err error
 	var t *template.Template
 	if !old {
@@ -45,6 +45,7 @@ func CreateInfra(namespace, registry, nmiVersion, micVersion, templateOutputPath
 		NMIArg              bool
 		EnableScaleFeatures bool
 		ImmutableUserMSIs   string
+		NMIMode             string
 	}{
 		namespace,
 		registry,
@@ -54,6 +55,7 @@ func CreateInfra(namespace, registry, nmiVersion, micVersion, templateOutputPath
 		nmiArg,
 		enableScaleFeatures,
 		immutableUserMSIs,
+		nmiMode,
 	}
 	if err := t.Execute(deployFile, deployData); err != nil {
 		return errors.Wrap(err, "Failed to create a deployment file from deployment-rbac.yaml")
diff --git a/test/e2e/README.md b/test/e2e/README.md
index 85facf83..9921868d 100644
--- a/test/e2e/README.md
+++ b/test/e2e/README.md
@@ -26,6 +26,9 @@ export KEYVAULT_SECRET_NAME='...'
 
 # The version of the secret inserted into the keyvault
 export KEYVAULT_SECRET_VERSION='...'
+
+# The client ID of the immutable user assigned MSI created when running setup.sh
+export IMMUTABLE_IDENTITY_CLIENT_ID='...'
 ```
 
 Optionally, to use custom images:
@@ -43,6 +46,8 @@ export MIC_VERSION='...'
 # The version of the identity validator to test. Defaults to `1.4`.
 export IDENTITY_VALIDATOR_VERSION='...'
 
+# The mode in which NMI is run. Allowed values: standard, managed. Default is standard
+export NMI_MODE='...'
 ```
 
 If you are using system asssigned identity cluster, please set the following variable:
diff --git a/test/e2e/aadpodidentity_test.go b/test/e2e/aadpodidentity_test.go
index a55d7c10..3cead7f7 100644
--- a/test/e2e/aadpodidentity_test.go
+++ b/test/e2e/aadpodidentity_test.go
@@ -61,6 +61,8 @@ var _ = BeforeSuite(func() {
 	Expect(err).NotTo(HaveOccurred())
 	cfg = *c
 	fmt.Printf("System MSI enabled: %v\n", cfg.SystemMSICluster)
+	fmt.Printf("NMI running in %s mode\n", cfg.NmiMode)
+	fmt.Printf("Running MIC version: %s, NMI Version: %s\n\n", cfg.MICVersion, cfg.NMIVersion)
 	setupInfra(cfg.Registry, cfg.NMIVersion, cfg.MICVersion, cfg.EnableScaleFeatures, cfg.ImmutableUserMSIs)
 })
 
@@ -127,7 +129,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 		_, err := cmd.CombinedOutput()
 		Expect(err).NotTo(HaveOccurred())
 
-		ok, err := azureassignedidentity.WaitOnLengthMatched(0)
+		ok, err := azureassignedidentity.WaitOnLengthMatched(0, false)
 		Expect(err).NotTo(HaveOccurred())
 		Expect(ok).To(Equal(true))
 
@@ -137,7 +139,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 	It("should pass the identity validating test", func() {
 		setUpIdentityAndDeployment(keyvaultIdentity, "", "1")
 
-		ok, err := azureassignedidentity.WaitOnLengthMatched(1)
+		ok, err := azureassignedidentity.WaitOnLengthMatched(1, true)
 		Expect(err).NotTo(HaveOccurred())
 		Expect(ok).To(Equal(true))
 
@@ -153,7 +155,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 		err := azureidentity.DeleteOnCluster(keyvaultIdentity, templateOutputPath)
 		Expect(err).NotTo(HaveOccurred())
 
-		ok, err := azureassignedidentity.WaitOnLengthMatched(0)
+		ok, err := azureassignedidentity.WaitOnLengthMatched(0, false)
 		Expect(err).NotTo(HaveOccurred())
 		Expect(ok).To(Equal(true))
 
@@ -175,7 +177,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 		err := azureidentitybinding.Delete(keyvaultIdentity, templateOutputPath)
 		Expect(err).NotTo(HaveOccurred())
 
-		ok, err := azureassignedidentity.WaitOnLengthMatched(0)
+		ok, err := azureassignedidentity.WaitOnLengthMatched(0, false)
 		Expect(err).NotTo(HaveOccurred())
 		Expect(ok).To(Equal(true))
 
@@ -196,7 +198,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 
 		waitForDeployDeletion(identityValidator)
 
-		ok, err := azureassignedidentity.WaitOnLengthMatched(0)
+		ok, err := azureassignedidentity.WaitOnLengthMatched(0, false)
 		Expect(err).NotTo(HaveOccurred())
 		Expect(ok).To(Equal(true))
 	})
@@ -216,7 +218,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 		// Drain the node that contains identity validator
 		node.Drain(nodeName)
 
-		ok, err := azureassignedidentity.WaitOnLengthMatched(1)
+		ok, err := azureassignedidentity.WaitOnLengthMatched(1, true)
 		Expect(ok).To(Equal(true))
 		Expect(err).NotTo(HaveOccurred())
 
@@ -243,7 +245,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 
 			setUpIdentityAndDeployment(keyvaultIdentity, fmt.Sprintf("%d", i), "1")
 
-			ok, err := azureassignedidentity.WaitOnLengthMatched(i + 1)
+			ok, err := azureassignedidentity.WaitOnLengthMatched(i+1, true)
 			Expect(ok).To(Equal(true))
 			Expect(err).NotTo(HaveOccurred())
 
@@ -280,7 +282,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 			err := azureidentity.DeleteOnCluster(data.identityName, templateOutputPath)
 			Expect(err).NotTo(HaveOccurred())
 
-			ok, err := azureassignedidentity.WaitOnLengthMatched(5 - 1 - i)
+			ok, err := azureassignedidentity.WaitOnLengthMatched(5-1-i, true)
 			Expect(err).NotTo(HaveOccurred())
 			Expect(ok).To(Equal(true))
 
@@ -309,7 +311,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 
 		setUpIdentityAndDeployment(keyvaultIdentity, "", "1")
 
-		ok, err := azureassignedidentity.WaitOnLengthMatched(1)
+		ok, err := azureassignedidentity.WaitOnLengthMatched(1, true)
 		Expect(err).NotTo(HaveOccurred())
 		Expect(ok).To(Equal(true))
 
@@ -344,12 +346,12 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 
 		// Delete pod identity to verify that the VM identity did not get deleted
 		waitForDeployDeletion(identityValidator)
-		cmd := exec.Command("kubectl", "delete", "AzureIdentity,AzureIdentityBinding,AzureAssignedIdentity", "--all")
+		cmd := exec.Command("kubectl", "delete", "AzureIdentity,AzureIdentityBinding", "--all")
 		util.PrintCommand(cmd)
 		_, err = cmd.CombinedOutput()
 		Expect(err).NotTo(HaveOccurred())
 
-		ok, err = azureassignedidentity.WaitOnLengthMatched(0)
+		ok, err = azureassignedidentity.WaitOnLengthMatched(0, false)
 		Expect(err).NotTo(HaveOccurred())
 		Expect(ok).To(Equal(true))
 
@@ -431,7 +433,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 		validateAzureAssignedIdentity(azureAssignedIdentity, keyvaultIdentity)
 
 		waitForDeployDeletion(identityValidator)
-		ok, err := azureassignedidentity.WaitOnLengthMatched(0)
+		ok, err := azureassignedidentity.WaitOnLengthMatched(0, false)
 		Expect(err).NotTo(HaveOccurred())
 		Expect(ok).To(Equal(true))
 
@@ -454,7 +456,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 		}
 
 		// WaitOnLengthMatched waits for 2 mins 30 seconds, so this will ensure we are performant at high scale
-		ok, err := azureassignedidentity.WaitOnLengthMatched(40)
+		ok, err := azureassignedidentity.WaitOnLengthMatched(40, true)
 		Expect(ok).To(Equal(true))
 		Expect(err).NotTo(HaveOccurred())
 
@@ -492,7 +494,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 
 		setUpIdentityAndDeployment(keyvaultIdentity, "", "1")
 
-		ok, err = azureassignedidentity.WaitOnLengthMatched(1)
+		ok, err = azureassignedidentity.WaitOnLengthMatched(1, false)
 		Expect(err).NotTo(HaveOccurred())
 		Expect(ok).To(Equal(true))
 
@@ -585,7 +587,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 			d.NodeName = vmss[0].Name
 		})
 
-		ok, err := azureassignedidentity.WaitOnLengthMatched(1)
+		ok, err := azureassignedidentity.WaitOnLengthMatched(1, true)
 		Expect(err).NotTo(HaveOccurred())
 		Expect(ok).To(Equal(true))
 
@@ -597,7 +599,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 
 		waitForDeployDeletion(identityValidator)
 
-		ok, err = azureassignedidentity.WaitOnLengthMatched(0)
+		ok, err = azureassignedidentity.WaitOnLengthMatched(0, false)
 		Expect(err).NotTo(HaveOccurred())
 		Expect(ok).To(Equal(true))
 
@@ -691,7 +693,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 
 		setUpIdentityAndDeployment(keyvaultIdentity, "0", "1")
 
-		ok, err = azureassignedidentity.WaitOnLengthMatched(1)
+		ok, err = azureassignedidentity.WaitOnLengthMatched(1, true)
 		Expect(err).NotTo(HaveOccurred())
 		Expect(ok).To(Equal(true))
 
@@ -749,7 +751,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 		err = azureidentitybinding.Create(fmt.Sprintf("%s-%d", keyvaultIdentity, 5), keyvaultIdentity, templateOutputPath)
 		Expect(err).NotTo(HaveOccurred())
 
-		ok, err := azureassignedidentity.WaitOnLengthMatched(2)
+		ok, err := azureassignedidentity.WaitOnLengthMatched(2, true)
 		Expect(err).NotTo(HaveOccurred())
 		Expect(ok).To(Equal(true))
 
@@ -773,7 +775,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 	It("should delete assigned identity when identity no longer exists on underlying node", func() {
 		setUpIdentityAndDeployment(keyvaultIdentity, "", "1")
 
-		ok, err := azureassignedidentity.WaitOnLengthMatched(1)
+		ok, err := azureassignedidentity.WaitOnLengthMatched(1, true)
 		Expect(err).NotTo(HaveOccurred())
 		Expect(ok).To(Equal(true))
 
@@ -788,7 +790,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 		removeUserAssignedIdentityFromCluster(nodeList, keyvaultIdentity)
 
 		waitForDeployDeletion(identityValidator)
-		ok, err = azureassignedidentity.WaitOnLengthMatched(0)
+		ok, err = azureassignedidentity.WaitOnLengthMatched(0, false)
 		Expect(err).NotTo(HaveOccurred())
 		Expect(ok).To(Equal(true))
 	})
@@ -825,7 +827,7 @@ var _ = Describe("Kubernetes cluster using aad-pod-identity", func() {
 			d.InitContainer = true
 		})
 
-		ok, err := azureassignedidentity.WaitOnLengthMatched(1)
+		ok, err := azureassignedidentity.WaitOnLengthMatched(1, true)
 		Expect(err).NotTo(HaveOccurred())
 		Expect(ok).To(Equal(true))
 
@@ -909,7 +911,7 @@ func runValidatorTest(iterations int) {
 		Expect(err).NotTo(HaveOccurred())
 		Expect(ok).To(Equal(true))
 
-		ok, err = azureassignedidentity.WaitOnLengthMatched(1)
+		ok, err = azureassignedidentity.WaitOnLengthMatched(1, true)
 		Expect(err).NotTo(HaveOccurred())
 		Expect(ok).To(Equal(true))
 
@@ -919,7 +921,7 @@ func runValidatorTest(iterations int) {
 
 		deleteAllIdentityValidator()
 
-		ok, err = azureassignedidentity.WaitOnLengthMatched(0)
+		ok, err = azureassignedidentity.WaitOnLengthMatched(0, false)
 		Expect(err).NotTo(HaveOccurred())
 		Expect(ok).To(Equal(true))
 	}
@@ -1105,7 +1107,7 @@ func checkInfra() {
 // setupInfra creates the crds, mic, nmi and blocks until iptable entries exist
 func setupInfraOld(registry, nmiVersion, micVersion string, immutableUserMSIs string) {
 	// Install CRDs and deploy MIC and NMI
-	err := infra.CreateInfra("default", registry, nmiVersion, micVersion, templateOutputPath, true, false, immutableUserMSIs)
+	err := infra.CreateInfra("default", registry, nmiVersion, micVersion, templateOutputPath, true, false, immutableUserMSIs, cfg.NmiMode)
 	Expect(err).NotTo(HaveOccurred())
 	checkInfra()
 }
@@ -1113,7 +1115,7 @@ func setupInfraOld(registry, nmiVersion, micVersion string, immutableUserMSIs st
 // setupInfra creates the crds, mic, nmi and blocks until iptable entries exist
 func setupInfra(registry, nmiVersion, micVersion string, enableScaleFeatures bool, immutableUserMSIs string) {
 	// Install CRDs and deploy MIC and NMI
-	err := infra.CreateInfra("default", registry, nmiVersion, micVersion, templateOutputPath, false, enableScaleFeatures, immutableUserMSIs)
+	err := infra.CreateInfra("default", registry, nmiVersion, micVersion, templateOutputPath, false, enableScaleFeatures, immutableUserMSIs, cfg.NmiMode)
 	Expect(err).NotTo(HaveOccurred())
 	checkInfra()
 }
@@ -1204,7 +1206,7 @@ func waitForDeployDeletion(deployName string) {
 }
 
 func deleteAllIdentityValidator() {
-	cmd := exec.Command("kubectl", "delete", "deploy", "-l", "app="+identityValidator)
+	cmd := exec.Command("kubectl", "delete", "deploy", "-l", "app="+identityValidator, "--force", "--grace-period", "0")
 	util.PrintCommand(cmd)
 	_, err := cmd.CombinedOutput()
 	Expect(err).NotTo(HaveOccurred())
diff --git a/test/e2e/config/config.go b/test/e2e/config/config.go
index ae3a692a..9af005a3 100644
--- a/test/e2e/config/config.go
+++ b/test/e2e/config/config.go
@@ -21,6 +21,7 @@ type Config struct {
 	SystemMSICluster         bool   `envconfig:"SYSTEM_MSI_CLUSTER" default:"false"`
 	EnableScaleFeatures      bool   `envconfig:"ENABLE_SCALE_FEATURES" default:"false"`
 	ImmutableUserMSIs        string `envconfig:"IMMUTABLE_IDENTITY_CLIENT_ID"`
+	NmiMode                  string `envconfig:"NMI_MODE" default:"standard"`
 }
 
 // ParseConfig will parse needed environment variables for running the tests
diff --git a/test/e2e/template/deployment-rbac.yaml b/test/e2e/template/deployment-rbac.yaml
index b8e091a7..619a830e 100644
--- a/test/e2e/template/deployment-rbac.yaml
+++ b/test/e2e/template/deployment-rbac.yaml
@@ -119,6 +119,7 @@ spec:
           {{if .NMIArg}}- nmi {{ end }}
           - "--host-ip=$(HOST_IP)"
           - "--node=$(NODE_NAME)"
+          - "--operation-mode={{.NMIMode}}"
           {{if .EnableScaleFeatures}}- "--enableScaleFeatures=true" {{end}}
         env:
           - name: HOST_IP