Switch back to 3 classes, with optional param this time

Author: rayluo

msal/__init__.py

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

msal/imds.py

@@ -41,55 +41,59 @@ class ManagedIdentity(UserDict):
     }
 
     @classmethod
-    def system_assigned(cls):
-        """Construct a system-assigned managed identity.
-
-        The outcome is equivalent to::
-
-            {"ManagedIdentityIdType": "SystemAssigned", "Id": None}
-        """
-        return ManagedIdentity(id_type=cls.SYSTEM_ASSIGNED)
+    def is_managed_identity(cls, unknown):
+        return isinstance(unknown, dict) and cls.ID_TYPE in unknown
 
     @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.
+    def is_user_assigned(cls, unknown):
+        return (
+            isinstance(unknown, dict)
+            and unknown.get(cls.ID_TYPE) in cls._types_mapping
+            and unknown.get(cls.ID))
 
-        The outcome is equivalent to::
+    def __init__(self, identifier=None, id_type=None):
+        # Undocumented. Use subclasses instead.
+        super(ManagedIdentity, self).__init__({
+            self.ID_TYPE: id_type,
+            self.ID: identifier,
+        })
 
-            {"ManagedIdentityIdType": "ClientId", "Id": "foo"}
-        """
-        return ManagedIdentity(identifier=identifier, id_type=cls.CLIENT_ID)
 
-    @classmethod
-    def user_assigned_resource_id(cls, identifier):
-        """Construct a ``ManagedIdentity`` instance from a user-assigned resource id.
+class SystemAssignedManagedIdentity(ManagedIdentity):
+    """Construct a system-assigned managed identity, which is equivalent to:
+    ``{"ManagedIdentityIdType": "SystemAssigned", "Id": None}``
+    """
+    def __init__(self):
+        super(SystemAssignedManagedIdentity, self).__init__(id_type=self.SYSTEM_ASSIGNED)
 
-        The outcome is equivalent to::
 
-            {"ManagedIdentityIdType": "ResourceId", "Id": "foo"}
-        """
-        return ManagedIdentity(identifier=identifier, id_type=cls.RESOURCE_ID)
-
-    @classmethod
-    def user_assigned_object_id(cls, identifier):
-        """Construct a ManagedIdentity instance from a user-assigned object id.
+class UserAssignedManagedIdentity(ManagedIdentity):
+    def __init__(self, client_id=None, resource_id=None, object_id=None):
+        """Construct a user-assigned managed identity.
 
-        The outcome will be equivalent to::
+        Depends on the id you provided, the outcome is equivalent to one of below::
 
+            {"ManagedIdentityIdType": "ClientId", "Id": "foo"}
+            {"ManagedIdentityIdType": "ResourceId", "Id": "foo"}
             {"ManagedIdentityIdType": "ObjectId", "Id": "foo"}
         """
-        return ManagedIdentity(identifier=identifier, id_type=cls.OBJECT_ID)
-
-    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,
-        })
+        if client_id and not resource_id and not object_id:
+            super(UserAssignedManagedIdentity, self).__init__(
+                id_type=self.CLIENT_ID, identifier=client_id)
+        elif not client_id and resource_id and not object_id:
+            super(UserAssignedManagedIdentity, self).__init__(
+                id_type=self.RESOURCE_ID, identifier=resource_id)
+        elif not client_id and not resource_id and object_id:
+            super(UserAssignedManagedIdentity, self).__init__(
+                id_type=self.OBJECT_ID, identifier=object_id)
+        else:
+            raise ValueError(
+                "You shall specify one of the three parameters: "
+                "client_id, resource_id, object_id")
 
 
 def _scope_to_resource(scope):  # This is an experimental reasonable-effort approach
@@ -268,7 +272,7 @@ def __init__(self, http_client, managed_identity, token_cache=None):
             import msal, requests
             client = msal.ManagedIdentityClient(
                 requests.Session(),
-                msal.UserAssignedManagedIdentity.from_client_id("foo"),
+                msal.UserAssignedManagedIdentity(client_id="foo"),
                 )
 
         Recipe: Write once, run everywhere.

tests/test_mi.py

@@ -9,22 +9,29 @@
 import requests
 
 from tests.http_client import MinimalResponse
-from msal import TokenCache, ManagedIdentity, ManagedIdentityClient
+from msal import (
+    TokenCache,
+    SystemAssignedManagedIdentity, UserAssignedManagedIdentity,
+    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(
-            ManagedIdentity.user_assigned_client_id("foo"),
+            UserAssignedManagedIdentity(client_id="foo"),
             {"ManagedIdentityIdType": "ClientId", "Id": "foo"})
         self.assertEqual(
-            ManagedIdentity.user_assigned_resource_id("foo"),
+            UserAssignedManagedIdentity(resource_id="foo"),
             {"ManagedIdentityIdType": "ResourceId", "Id": "foo"})
         self.assertEqual(
-            ManagedIdentity.user_assigned_object_id("foo"),
+            UserAssignedManagedIdentity(object_id="foo"),
             {"ManagedIdentityIdType": "ObjectId", "Id": "foo"})
+        with self.assertRaises(ValueError):
+            UserAssignedManagedIdentity()
+        with self.assertRaises(ValueError):
+            UserAssignedManagedIdentity(client_id="foo", resource_id="bar")
         self.assertEqual(
-            ManagedIdentity.system_assigned(),
+            SystemAssignedManagedIdentity(),
             {"ManagedIdentityIdType": "SystemAssigned", "Id": None})