[Core] aaz: Support managed identity

src/azure-cli-core/azure/cli/core/aaz/__init__.py

@@ -22,7 +22,7 @@
 from ._command import AAZCommand, AAZWaitCommand, AAZCommandGroup, \
     register_callback, register_command, register_command_group, load_aaz_command_table, link_helper
 from ._field_type import AAZIntType, AAZFloatType, AAZStrType, AAZBoolType, AAZDictType, AAZFreeFormDictType, \
-    AAZListType, AAZObjectType
+    AAZListType, AAZObjectType, AAZIdentityObjectType
 from ._operation import AAZHttpOperation, AAZJsonInstanceUpdateOperation, AAZGenericInstanceUpdateOperation, \
     AAZJsonInstanceDeleteOperation, AAZJsonInstanceCreateOperation
 from ._prompt import AAZPromptInput, AAZPromptPasswordInput

src/azure-cli-core/azure/cli/core/aaz/_field_type.py

@@ -6,7 +6,8 @@
 
 from collections import OrderedDict
 from ._base import AAZBaseType, AAZValuePatch, AAZUndefined
-from ._field_value import AAZObject, AAZDict, AAZFreeFormDict, AAZList, AAZSimpleValue
+from ._field_value import AAZObject, AAZDict, AAZFreeFormDict, AAZList, AAZSimpleValue, \
+    AAZIdentityObject
 from ._utils import to_snack_case
 from .exceptions import AAZUnknownFieldError, AAZConflictFieldDefinitionError, AAZValuePrecisionLossError, \
     AAZInvalidFieldError, AAZInvalidValueError
@@ -409,3 +410,7 @@ def process_data(self, data, **kwargs):
                 value[idx] = sub_data
 
         return result
+
+
+class AAZIdentityObjectType(AAZObjectType):
+    _ValueCls = AAZIdentityObject

src/azure-cli-core/azure/cli/core/aaz/_field_value.py

@@ -2,6 +2,7 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
+# pylint: disable=protected-access
 import copy
 
 from ._base import AAZBaseValue, AAZValuePatch, AAZUndefined
@@ -444,3 +445,93 @@ def to_serialized_data(self, processor=None, keep_undefined_in_list=False,  # py
         if processor:
             result = processor(self._schema, result)
         return result
+
+
+class AAZIdentityObject(AAZObject):  # pylint: disable=too-few-public-methods
+    def to_serialized_data(self, processor=None, **kwargs):
+        calculate_data = dict()
+        if self._data == AAZUndefined:
+            result = AAZUndefined
+
+        elif self._data is None:
+            result = None
+
+        else:
+            result = {}
+            schema = self._schema
+
+            for name, field_schema in schema._fields.items():
+                if name in self._data:
+                    v = self[name].to_serialized_data(processor=processor, **kwargs)
+                    if v == AAZUndefined:
+                        continue
+
+                    if field_schema._serialized_name:
+                        name = field_schema._serialized_name
+
+                    if name in {"userAssigned", "systemAssigned"}:
+                        calculate_data[name] = v
+                        calculate_data["action"] = field_schema._flags.get("action", None)  # no action in GET operation
+
+                    else:
+                        result[name] = v
+
+        result = self._build_identity(calculate_data, result)
+
+        if not result:
+            result = {"type": "None"}  # empty identity
+
+        if not result and self._is_patch:
+            result = AAZUndefined
+
+        if processor:
+            result = processor(self._schema, result)
+
+        return result
+
+    def _build_identity(self, calculate_data, result):
+        action = calculate_data.get("action", None)
+        if not action:
+            return result
+
+        user_assigned = calculate_data.get("userAssigned", None)
+        system_assigned = calculate_data.get("systemAssigned", None)
+
+        identities = set(result.pop("userAssignedIdentities", {}).keys())
+        has_system_identity = "systemassigned" in result.pop("type", "").lower()
+
+        if action == "remove":
+            if user_assigned is not None:
+                if len(user_assigned) > 1:  # remove each
+                    identities -= set(user_assigned)
+
+                else:  # remove all
+                    identities = {}
+
+            if identities:
+                result["userAssignedIdentities"] = {k: {} for k in identities}
+                if system_assigned or not has_system_identity:
+                    result["type"] = "UserAssigned"
+
+                else:
+                    result["type"] = "SystemAssigned,UserAssigned"
+
+            elif not system_assigned and has_system_identity:
+                result["type"] = "SystemAssigned"
+
+        else:  # assign or create
+            if user_assigned:
+                identities |= set(user_assigned)
+
+            if identities:
+                result["userAssignedIdentities"] = {k: {} for k in identities}
+                if system_assigned or has_system_identity:
+                    result["type"] = "SystemAssigned,UserAssigned"
+
+                else:
+                    result["type"] = "UserAssigned"
+
+            elif system_assigned or has_system_identity:
+                result["type"] = "SystemAssigned"
+
+        return result

src/azure-cli-core/azure/cli/core/tests/test_aaz_identity.py

@@ -0,0 +1,152 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+import unittest
+
+from azure.cli.core.aaz import AAZIdentityObjectType
+from azure.cli.core.aaz._field_value import AAZIdentityObject
+
+
+class TestAAZIdentity(unittest.TestCase):
+    identity = AAZIdentityObject(schema=AAZIdentityObjectType(), data={})
+
+    def test_main_command(self):
+        result = {}
+        data = {
+            "userAssigned": ["a", "b"],
+            "systemAssigned": "SystemAssigned",
+            "action": "create",
+        }
+        result = self.identity._build_identity(data, result)
+        self.assertTrue(result["type"] == "SystemAssigned,UserAssigned")
+        self.assertTrue(result["userAssignedIdentities"] == {"a": {}, "b": {}})
+
+        result = {}
+        data = {
+            "userAssigned": [],
+            "systemAssigned": "SystemAssigned",
+            "action": "create",
+        }
+        result = self.identity._build_identity(data, result)
+        self.assertTrue(result["type"] == "SystemAssigned")
+        self.assertTrue("userAssignedIdentities" not in result)
+
+        result = {}
+        data = {
+            "userAssigned": [],
+            "action": "create",
+        }
+        result = self.identity._build_identity(data, result)
+        self.assertTrue(result == {})
+
+        result = {}
+        data = {"action": "create"}
+        result = self.identity._build_identity(data, result)
+        self.assertTrue(result == {})
+
+    def test_assign_command(self):
+        result = {
+            "type": "systemAssigned, userAssigned",
+            "userAssignedIdentities": {"a": {}, "b": {}},
+        }
+        data = {
+            "userAssigned": ["b", "c"],
+            "action": "assign",
+        }
+        result = self.identity._build_identity(data, result)
+        self.assertTrue(result["type"] == "SystemAssigned,UserAssigned")
+        self.assertTrue(result["userAssignedIdentities"] == {"a": {}, "b": {}, "c": {}})
+
+        result = {
+            "type": "userAssigned",
+            "userAssignedIdentities": {"a": {}, "b": {}},
+        }
+        data = {
+            "systemAssigned": "SystemAssigned",
+            "action": "assign",
+        }
+        result = self.identity._build_identity(data, result)
+        self.assertTrue(result["type"] == "SystemAssigned,UserAssigned")
+        self.assertTrue(result["userAssignedIdentities"] == {"a": {}, "b": {}})
+
+        result = {
+            "type": "systemAssigned",
+        }
+        data = {
+            "userAssigned": ["a", "b"],
+            "action": "assign",
+        }
+        result = self.identity._build_identity(data, result)
+        self.assertTrue(result["type"] == "SystemAssigned,UserAssigned")
+        self.assertTrue(result["userAssignedIdentities"] == {"a": {}, "b": {}})
+
+        result = {}
+        data = {
+            "userAssigned": ["a", "b"],
+            "action": "assign",
+        }
+        result = self.identity._build_identity(data, result)
+        self.assertTrue(result["type"] == "UserAssigned")
+        self.assertTrue(result["userAssignedIdentities"] == {"a": {}, "b": {}})
+
+        result = {}
+        data = {
+            "userAssigned": [],
+            "action": "assign",
+        }
+        result = self.identity._build_identity(data, result)
+        self.assertTrue(result == {})
+
+    def test_remove_command(self):
+        result = {
+            "type": "systemAssigned, userAssigned",
+            "userAssignedIdentities": {"a": {}, "b": {}},
+        }
+        data = {
+            "userAssigned": ["b", "c"],
+            "systemAssigned": "SystemAssigned",
+            "action": "remove",
+        }
+        result = self.identity._build_identity(data, result)
+        self.assertTrue(result["type"] == "UserAssigned")
+        self.assertTrue(result["userAssignedIdentities"] == {"a": {}})
+
+        result = {
+            "type": "systemAssigned, userAssigned",
+            "userAssignedIdentities": {"a": {}, "b": {}},
+        }
+        data = {
+            "userAssigned": [],
+            "action": "remove",
+        }
+        result = self.identity._build_identity(data, result)
+        self.assertTrue(result["type"] == "SystemAssigned")
+        self.assertTrue("userAssignedIdentities" not in result)
+
+        result = {
+            "type": "userAssigned",
+            "userAssignedIdentities": {"a": {}, "b": {}},
+        }
+        data = {
+            "userAssigned": [],
+            "action": "remove",
+        }
+        result = self.identity._build_identity(data, result)
+        self.assertTrue(result == {})
+
+        result = {}
+        data = {
+            "userAssigned": ["a", "b"],
+            "action": "remove",
+        }
+        result = self.identity._build_identity(data, result)
+        self.assertTrue(result == {})
+
+        result = {}
+        data = {
+            "userAssigned": [],
+            "action": "remove",
+        }
+        result = self.identity._build_identity(data, result)
+        self.assertTrue(result == {})