Simplify 3 managed identity classes into 1

Author: rayluo

msal/__init__.py

@@ -33,9 +33,5 @@
     )
 from .oauth2cli.oidc import Prompt
 from .token_cache import TokenCache, SerializableTokenCache
-from .imds import (
-    SystemAssignedManagedIdentity,
-    UserAssignedManagedIdentity,
-    ManagedIdentityClient,
-    )
+from .imds import ManagedIdentity, ManagedIdentityClient
 

msal/imds.py

@@ -21,83 +21,75 @@
 logger = logging.getLogger(__name__)
 
 class ManagedIdentity(UserDict):
+    """Feed an instance of this class to :class:`msal.ManagedIdentityClient`
+    to acquire token for the specified managed identity.
+    """
     # The key names used in config dict
-    ID_TYPE = "ManagedIdentityIdType"
+    ID_TYPE = "ManagedIdentityIdType"  # Contains keyword ManagedIdentity so its json equivalent will be more readable
     ID = "Id"
-    def __init__(self, identifier=None, id_type=None):
-        super(ManagedIdentity, self).__init__({
-            self.ID_TYPE: id_type,
-            self.ID: identifier,
-        })
-
 
-class UserAssignedManagedIdentity(ManagedIdentity):
-    """Feed an instance of this class to :class:`msal.ManagedIdentityClient`
-    to acquire token for user-assigned managed identity.
-    """
+    # Valid values for key ID_TYPE
     CLIENT_ID = "ClientId"
     RESOURCE_ID = "ResourceId"
     OBJECT_ID = "ObjectId"
+    SYSTEM_ASSIGNED = "SystemAssigned"
+
     _types_mapping = {  # Maps type name in configuration to type name on wire
         CLIENT_ID: "client_id",
         RESOURCE_ID: "mi_res_id",
         OBJECT_ID: "object_id",
     }
-    def __init__(self, identifier, id_type):
-        """Do not use this contructor. Use the following factory methods instead."""
-        if id_type not in self._types_mapping:
-            raise ValueError("id_type only accepts one of: {}".format(
-                list(self._types_mapping)))
-        super(UserAssignedManagedIdentity, self).__init__(
-            identifier=identifier,
-            id_type=id_type,
-        )
 
     @classmethod
-    def from_client_id(cls, identifier):
-        """Construct a UserAssignedManagedIdentity instance from a client id.
+    def system_assigned(cls):
+        """Construct a system-assigned managed identity.
 
-        The outcome will be equivalent to::
+        The outcome is equivalent to::
+
+            {"ManagedIdentityIdType": "SystemAssigned", "Id": None}
+        """
+        return ManagedIdentity(id_type=cls.SYSTEM_ASSIGNED)
+
+    @classmethod
+    def is_system_assigned(cls, unknown):
+        return isinstance(unknown, dict) and unknown.get(cls.ID_TYPE) == cls.SYSTEM_ASSIGNED
+
+    @classmethod
+    def user_assigned_client_id(cls, identifier):
+        """Construct a ``ManagedIdentity`` instance from a user-assigned client id.
+
+        The outcome is equivalent to::
 
             {"ManagedIdentityIdType": "ClientId", "Id": "foo"}
         """
-        return UserAssignedManagedIdentity(identifier, cls.CLIENT_ID)
+        return ManagedIdentity(identifier=identifier, id_type=cls.CLIENT_ID)
 
     @classmethod
-    def from_resource_id(cls, identifier):
-        """Construct a UserAssignedManagedIdentity instance from a resource id.
+    def user_assigned_resource_id(cls, identifier):
+        """Construct a ``ManagedIdentity`` instance from a user-assigned resource id.
 
-        The outcome will be equivalent to::
+        The outcome is equivalent to::
 
             {"ManagedIdentityIdType": "ResourceId", "Id": "foo"}
         """
-        return UserAssignedManagedIdentity(identifier, cls.RESOURCE_ID)
+        return ManagedIdentity(identifier=identifier, id_type=cls.RESOURCE_ID)
 
     @classmethod
-    def from_object_id(cls, identifier):
-        """Construct a UserAssignedManagedIdentity instance from an object id.
+    def user_assigned_object_id(cls, identifier):
+        """Construct a ManagedIdentity instance from a user-assigned object id.
 
         The outcome will be equivalent to::
 
             {"ManagedIdentityIdType": "ObjectId", "Id": "foo"}
         """
-        return UserAssignedManagedIdentity(identifier, cls.OBJECT_ID)
-
-
-class SystemAssignedManagedIdentity(ManagedIdentity):
-    """Feed an instance of this class to :class:`msal.ManagedIdentityClient`
-    to acquire token for system-assigned managed identity.
+        return ManagedIdentity(identifier=identifier, id_type=cls.OBJECT_ID)
 
-    By design, an instance of this class is equivalent to::
-
-        {"ManagedIdentityIdType": "SystemAssignedManagedIdentity", "Id": None}
-    """
-    def __init__(self):
-        super(SystemAssignedManagedIdentity, self).__init__(
-            id_type="SystemAssignedManagedIdentity",  # As of this writing,
-                # It can be any value other than
-                # UserAssignedManagedIdentity._types_mapping's key names
-        )
+    def __init__(self, identifier=None, id_type=None):
+        # Undocumented. Use other class methods instead.
+        super(ManagedIdentity, self).__init__({
+            self.ID_TYPE: id_type,
+            self.ID: identifier,
+        })
 
 
 def _scope_to_resource(scope):  # This is an experimental reasonable-effort approach
@@ -136,7 +128,7 @@ def _obtain_token(http_client, managed_identity, resource):
 
 
 def _adjust_param(params, managed_identity):
-    id_name = UserAssignedManagedIdentity._types_mapping.get(
+    id_name = ManagedIdentity._types_mapping.get(
         managed_identity.get(ManagedIdentity.ID_TYPE))
     if id_name:
         params[id_name] = managed_identity[ManagedIdentity.ID]

tests/test_mi.py

@@ -9,25 +9,23 @@
 import requests
 
 from tests.http_client import MinimalResponse
-from msal import (
-    TokenCache,
-    SystemAssignedManagedIdentity, UserAssignedManagedIdentity, ManagedIdentityClient)
+from msal import TokenCache, ManagedIdentity, ManagedIdentityClient
 
 
 class ManagedIdentityTestCase(unittest.TestCase):
     def test_helper_class_should_be_interchangable_with_dict_which_could_be_loaded_from_file_or_env_var(self):
         self.assertEqual(
-            UserAssignedManagedIdentity.from_client_id("foo"),
+            ManagedIdentity.user_assigned_client_id("foo"),
             {"ManagedIdentityIdType": "ClientId", "Id": "foo"})
         self.assertEqual(
-            UserAssignedManagedIdentity.from_resource_id("foo"),
+            ManagedIdentity.user_assigned_resource_id("foo"),
             {"ManagedIdentityIdType": "ResourceId", "Id": "foo"})
         self.assertEqual(
-            UserAssignedManagedIdentity.from_object_id("foo"),
+            ManagedIdentity.user_assigned_object_id("foo"),
             {"ManagedIdentityIdType": "ObjectId", "Id": "foo"})
         self.assertEqual(
-            SystemAssignedManagedIdentity(),
-            {"ManagedIdentityIdType": "SystemAssignedManagedIdentity", "Id": None})
+            ManagedIdentity.system_assigned(),
+            {"ManagedIdentityIdType": "SystemAssigned", "Id": None})
 
 
 class ClientTestCase(unittest.TestCase):
@@ -38,7 +36,7 @@ def setUp(self):
             requests.Session(),
             {   # Here we test it with the raw dict form, to test that
                 # the client has no hard dependency on ManagedIdentity object
-                "ManagedIdentityIdType": "SystemAssignedManagedIdentity", "Id": None,
+                "ManagedIdentityIdType": "SystemAssigned", "Id": None,
             },
             token_cache=TokenCache())