[load] Adding CLI support for engine reference identity feature (#8492)

* updated sdk with 2024-12-01 version api

* Updating client stop call and response from list_metric_dimension_values

* Fixi for few regression issue

* Updated test recordings to use new API version for data plane

* Changes for azdev style

* initial params setup

* Added validators and support for config input

* Added test cases for engine mi feature

* Updated history and setup file

* Addressing PR comments

* Fixing naming of argument

* test recordings update

* Addressing PR comments

* Updating version to 1.6.0

Author: Himanshu49

src/load/HISTORY.rst

@@ -4,6 +4,11 @@ Release History
 ===============
 
 
+1.6.0
+++++++
+* Add support for engine reference identity using CLI. Engine reference identity can be set using `--engine-ref-id-type` and `--engine-ref-ids` argument in 'az load test create' and 'az load test update' commands. Engine reference identity set in YAML config file under key `referenceIdentities` with `kind` as `Engine` will also be honoured.
+
+
 1.5.0
 ++++++
 * Add support for Locust based load tests.

src/load/azext_load/data_plane/load_test/custom.py

@@ -56,6 +56,8 @@ def create_test(
     autostop_error_rate=None,
     autostop_error_rate_time_window=None,
     regionwise_engines=None,
+    engine_ref_id_type=None,
+    engine_ref_ids=None,
 ):
     client = get_admin_data_plane_client(cmd, load_test_resource, resource_group_name)
     logger.info("Create test has started for test ID : %s", test_id)
@@ -91,6 +93,8 @@ def create_test(
             disable_public_ip=disable_public_ip,
             autostop_criteria=autostop_criteria,
             regionwise_engines=regionwise_engines,
+            engine_ref_id_type=engine_ref_id_type,
+            engine_ref_ids=engine_ref_ids,
         )
     else:
         yaml = load_yaml(load_test_config_file)
@@ -119,6 +123,8 @@ def create_test(
             disable_public_ip=disable_public_ip,
             autostop_criteria=autostop_criteria,
             regionwise_engines=regionwise_engines,
+            engine_ref_id_type=engine_ref_id_type,
+            engine_ref_ids=engine_ref_ids,
         )
     logger.debug("Creating test with test ID: %s and body : %s", test_id, body)
     response = client.create_or_update_test(test_id=test_id, body=body)
@@ -158,6 +164,8 @@ def update_test(
     autostop_error_rate=None,
     autostop_error_rate_time_window=None,
     regionwise_engines=None,
+    engine_ref_id_type=None,
+    engine_ref_ids=None,
 ):
     client = get_admin_data_plane_client(cmd, load_test_resource, resource_group_name)
     logger.info("Update test has started for test ID : %s", test_id)
@@ -191,6 +199,8 @@ def update_test(
             disable_public_ip=disable_public_ip,
             autostop_criteria=autostop_criteria,
             regionwise_engines=regionwise_engines,
+            engine_ref_id_type=engine_ref_id_type,
+            engine_ref_ids=engine_ref_ids,
         )
     else:
         body = create_or_update_test_without_config(
@@ -208,6 +218,8 @@ def update_test(
             disable_public_ip=disable_public_ip,
             autostop_criteria=autostop_criteria,
             regionwise_engines=regionwise_engines,
+            engine_ref_id_type=engine_ref_id_type,
+            engine_ref_ids=engine_ref_ids
         )
     logger.info("Updating test with test ID: %s", test_id)
     response = client.create_or_update_test(test_id=test_id, body=body)

src/load/azext_load/data_plane/load_test/help.py

@@ -19,7 +19,7 @@
         az load test create --load-test-resource sample-alt-resource --resource-group sample-rg --test-id sample-test-id --load-test-config-file ~/resources/sample-config.yaml
     - name: Create a test with arguments.
       text: |
-        az load test create --test-id sample-test-id --load-test-resource sample-alt-resource --resource-group sample-rg --display-name "Sample Name" --description "Test description" --test-plan sample-jmx.jmx --engine-instances 1 --env rps=2 count=1
+        az load test create --test-id sample-test-id --load-test-resource sample-alt-resource --resource-group sample-rg --display-name "Sample Name" --description "Test description" --test-plan sample-jmx.jmx --engine-instances 1 --env rps=2 count=1 --engine-ref-id-type SystemAssigned
     - name: Create a test with load test config file and override engine-instance and env using arguments and don't wait for file upload.
       text: |
         az load test create --load-test-resource sample-alt-resource --resource-group sample-rg --test-id sample-test-id --load-test-config-file ~/resources/sample-config.yaml --engine-instances 1 --env rps=2 count=1 --no-wait
@@ -46,6 +46,9 @@
     - name: Create a Locust based load test
       text: |
         az load test create --test-id sample-test-id --load-test-resource sample-alt-resource --resource-group sample-rg --test-plan ~/resources/sample-locust-file.py --test-type Locust --env LOCUST_HOST="https://azure.microsoft.com" LOCUST_SPAWN_RATE=0.3 LOCUST_RUN_TIME=120 LOCUST_USERS=4
+    - name: Create a test with user assigned Managed Identity reference for engine.
+      text: |
+        az load test create --test-id sample-test-id --load-test-resource sample-alt-resource --resource-group sample-rg --display-name "Sample Name" --engine-ref-id-type UserAssigned --engine-ref-ids "/subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/sample-rg/providers/microsoft.managedidentity/userassignedidentities/sample-mi"
 """
 
 helps[
@@ -97,6 +100,9 @@
     - name: Update multi-region load configuration.
       text: |
         az load test update --load-test-resource sample-alt-resource --resource-group sample-rg --test-id sample-existing-test-id --engine-instances 5 --regionwise-engines eastus=2 westus2=1 eastasia=2
+    - name: Update a test with user assigned Managed Identity reference for engine to SystemAssigned.
+      text: |
+        az load test update --test-id sample-test-id --load-test-resource sample-alt-resource --resource-group sample-rg --display-name "Sample Name" --engine-ref-id-type SystemAssigned
 """
 
 helps[

src/load/azext_load/data_plane/load_test/params.py

@@ -34,6 +34,8 @@ def load_arguments(self, _):
         c.argument("autostop_error_rate", argtypes.autostop_error_rate)
         c.argument("autostop_error_rate_time_window", argtypes.autostop_error_rate_time_window)
         c.argument("regionwise_engines", argtypes.regionwise_engines)
+        c.argument("engine_ref_id_type", argtypes.engine_ref_id_type)
+        c.argument("engine_ref_ids", argtypes.engine_ref_ids)
 
     with self.argument_context("load test update") as c:
         c.argument("load_test_config_file", argtypes.load_test_config_file)
@@ -55,6 +57,8 @@ def load_arguments(self, _):
         c.argument("autostop_error_rate", argtypes.autostop_error_rate)
         c.argument("autostop_error_rate_time_window", argtypes.autostop_error_rate_time_window)
         c.argument("regionwise_engines", argtypes.regionwise_engines)
+        c.argument("engine_ref_id_type", argtypes.engine_ref_id_type)
+        c.argument("engine_ref_ids", argtypes.engine_ref_ids)
 
     with self.argument_context("load test set-baseline") as c:
         c.argument("test_run_id", argtypes.test_run_id)

src/load/azext_load/data_plane/utils/argtypes.py

@@ -392,6 +392,23 @@
     help="Specify the engine count for each region in the format: region1=engineCount1 region2=engineCount2 .... Use region names in the format accepted by Azure Resource Manager (ARM). Ensure the regions are supported by Azure Load Testing. Multi-region load tests can only target public endpoints.",
 )
 
+engine_ref_id_type = CLIArgumentType(
+    options_list=["--engine-ref-id-type"],
+    type=str,
+    completer=get_generic_completion_list(
+        utils.get_enum_values(models.EngineIdentityType)
+    ),
+    choices=utils.get_enum_values(models.EngineIdentityType),
+    help="Type of identity to be configured for the engine.",
+)
+
+engine_ref_ids = CLIArgumentType(
+    options_list=["--engine-ref-ids"],
+    nargs="+",
+    validator=validators.validate_engine_ref_ids,
+    help="Space separated list of fully qualified resource IDs of the managed identities to be configured on the engine. Required only for user assigned identities. ",
+)
+
 response_time_aggregate = CLIArgumentType(
     options_list=["--aggregation"],
     type=str,

src/load/azext_load/data_plane/utils/constants.py

@@ -28,6 +28,11 @@ class LoadTestConfigKeys:
     REGION = "region"
     QUICK_START = "quickStartTest"
     SPLIT_CSV = "splitAllCSVs"
+    REFERENCE_IDENTITIES = "referenceIdentities"
+    ENGINE = "Engine"
+    TYPE = "type"
+    KIND = "kind"
+    VALUE = "value"
 
 
 @dataclass

src/load/azext_load/data_plane/utils/models.py

@@ -11,6 +11,12 @@ class IdentityType(str, Enum):
     UserAssigned = "UserAssigned"
 
 
+class EngineIdentityType(str, Enum):
+    SystemAssigned = "SystemAssigned"
+    UserAssigned = "UserAssigned"
+    NoneValue = "None"
+
+
 class AllowedFileTypes(str, Enum):
     ADDITIONAL_ARTIFACTS = "ADDITIONAL_ARTIFACTS"
     JMX_FILE = "JMX_FILE"

src/load/azext_load/data_plane/utils/utils.py

@@ -20,7 +20,7 @@
 from azure.mgmt.core.tools import is_valid_resource_id, parse_resource_id
 from knack.log import get_logger
 
-from .models import IdentityType, AllowedFileTypes, AllowedTestTypes, AllowedTestPlanFileExtensions
+from .models import IdentityType, AllowedFileTypes, AllowedTestTypes, EngineIdentityType, AllowedTestPlanFileExtensions
 
 logger = get_logger(__name__)
 
@@ -306,6 +306,8 @@ def load_yaml(file_path):
         ) from e
 
 
+# pylint: disable=line-too-long
+# Disabling this because dictionary key are too long
 def convert_yaml_to_test(cmd, data):
     new_body = {}
     if LoadTestConfigKeys.DISPLAY_NAME in data:
@@ -337,8 +339,11 @@ def convert_yaml_to_test(cmd, data):
         new_body["passFailCriteria"] = utils_yaml_config.yaml_parse_failure_criteria(data=data)
     if data.get(LoadTestConfigKeys.AUTOSTOP) is not None:
         new_body["autoStopCriteria"] = utils_yaml_config.yaml_parse_autostop_criteria(data=data)
+
+    utils_yaml_config.update_engine_reference_identity(new_body, data)
     logger.debug("Converted yaml to test body: %s", new_body)
     return new_body
+# pylint: enable=line-too-long
 
 
 # pylint: disable=too-many-branches
@@ -360,6 +365,8 @@ def create_or_update_test_with_config(
     disable_public_ip=None,
     autostop_criteria=None,
     regionwise_engines=None,
+    engine_ref_id_type=None,
+    engine_ref_ids=None,
 ):
     logger.info(
         "Creating a request body for create or update test using config and parameters."
@@ -514,6 +521,20 @@ def create_or_update_test_with_config(
             "This can lead to incoming charges for an incorrectly configured test."
         )
 
+    # if argument is provided prefer that over yaml values
+    if engine_ref_id_type:
+        validators.validate_engine_ref_ids_and_type(engine_ref_id_type, engine_ref_ids)
+        if engine_ref_id_type:
+            new_body["engineBuiltinIdentityType"] = engine_ref_id_type
+        if engine_ref_ids:
+            new_body["engineBuiltinIdentityIds"] = engine_ref_ids
+    elif yaml_test_body.get("engineBuiltinIdentityType"):
+        new_body["engineBuiltinIdentityType"] = yaml_test_body.get("engineBuiltinIdentityType")
+        new_body["engineBuiltinIdentityIds"] = yaml_test_body.get("engineBuiltinIdentityIds")
+    else:
+        new_body["engineBuiltinIdentityType"] = body.get("engineBuiltinIdentityType")
+        new_body["engineBuiltinIdentityIds"] = body.get("engineBuiltinIdentityIds")
+
     logger.debug("Request body for create or update test: %s", new_body)
     return new_body
 
@@ -537,6 +558,8 @@ def create_or_update_test_without_config(
     autostop_criteria=None,
     regionwise_engines=None,
     baseline_test_run_id=None,
+    engine_ref_id_type=None,
+    engine_ref_ids=None,
 ):
     logger.info(
         "Creating a request body for test using parameters and old test body (in case of update)."
@@ -640,6 +663,21 @@ def create_or_update_test_without_config(
         )
     new_body["baselineTestRunId"] = baseline_test_run_id if baseline_test_run_id else body.get("baselineTestRunId")
 
+    # pylint: disable=line-too-long
+    # Disabling this because dictionary key are too long
+    # raises error if engine_reference_identity_type and corresponding identities is not a valid combination
+    validators.validate_engine_ref_ids_and_type(engine_ref_id_type, engine_ref_ids, body.get("engineBuiltinIdentityType"))
+    if engine_ref_id_type:
+        new_body["engineBuiltinIdentityType"] = engine_ref_id_type
+        if engine_ref_ids:
+            new_body["engineBuiltinIdentityIds"] = engine_ref_ids
+    else:
+        new_body["engineBuiltinIdentityType"] = body.get("engineBuiltinIdentityType")
+        if engine_ref_ids and body.get("engineBuiltinIdentityType") != EngineIdentityType.UserAssigned:
+            raise InvalidArgumentValueError("Engine reference identities can only be provided when engine reference identity type is user assigned")
+        new_body["engineBuiltinIdentityIds"] = engine_ref_ids if engine_ref_ids else body.get("engineBuiltinIdentityIds")
+    # pylint: enable=line-too-long
+
     logger.debug("Request body for create or update test: %s", new_body)
     return new_body
 

src/load/azext_load/data_plane/utils/utils_yaml_config.py

@@ -7,6 +7,8 @@
 
 from azext_load.data_plane.utils.constants import LoadTestConfigKeys
 from azext_load.data_plane.utils import validators
+from azext_load.data_plane.utils.models import EngineIdentityType
+from azure.mgmt.core.tools import is_valid_resource_id
 from azure.cli.core.azclierror import (
     InvalidArgumentValueError,
 )
@@ -137,3 +139,43 @@ def yaml_parse_loadtest_configuration(cmd, data):
     if data.get(LoadTestConfigKeys.SPLIT_CSV) is not None:
         load_test_configuration["splitAllCSVs"] = _yaml_parse_splitcsv(data=data)
     return load_test_configuration
+
+
+# pylint: disable=line-too-long
+# Disabling this because dictionary key are too long
+def yaml_parse_engine_identities(data):
+    engine_identities = []
+    reference_type = None
+    reference_identities = data.get(LoadTestConfigKeys.REFERENCE_IDENTITIES)
+    for identity in reference_identities:
+        if identity and identity.get(LoadTestConfigKeys.KIND) == LoadTestConfigKeys.ENGINE:
+            curr_ref_type = identity.get(LoadTestConfigKeys.TYPE)
+            curr_ref_value = identity.get(LoadTestConfigKeys.VALUE)
+            if reference_type and curr_ref_type != reference_type:
+                raise InvalidArgumentValueError(
+                    "Engine identity type should be either None, SystemAssigned, or UserAssigned. A combination of identity types are not supported."
+                )
+            if curr_ref_type != EngineIdentityType.UserAssigned:
+                if curr_ref_value:
+                    raise InvalidArgumentValueError(
+                        "Reference identity value should be provided only for UserAssigned identity type."
+                    )
+            else:
+                if not is_valid_resource_id(curr_ref_value):
+                    raise InvalidArgumentValueError(
+                        "%s is not a valid resource id" % curr_ref_value
+                    )
+            engine_identities.append(curr_ref_value)
+            reference_type = curr_ref_type
+    return reference_type, engine_identities
+
+
+def update_engine_reference_identity(new_body, data):
+    if data.get(LoadTestConfigKeys.REFERENCE_IDENTITIES):
+        for identity in data[LoadTestConfigKeys.REFERENCE_IDENTITIES]:
+            if identity and identity.get(LoadTestConfigKeys.KIND) == LoadTestConfigKeys.ENGINE:
+                new_body["engineBuiltinIdentityType"], new_body["engineBuiltinIdentityIds"] = yaml_parse_engine_identities(data=data)
+                if new_body["engineBuiltinIdentityType"] in [EngineIdentityType.NoneValue, EngineIdentityType.SystemAssigned]:
+                    new_body.pop("engineBuiltinIdentityIds")
+                break
+# pylint: enable=line-too-long

src/load/azext_load/data_plane/utils/validators.py

@@ -21,6 +21,7 @@
     AllowedMetricNamespaces,
     AllowedTestTypes,
     AllowedTestPlanFileExtensions,
+    EngineIdentityType,
 )
 
 logger = get_logger(__name__)
@@ -522,3 +523,28 @@ def validate_regionwise_engines(cmd, namespace):
             )
         regionwise_engines.append({"region": key.strip().lower(), "engineInstances": value})
     namespace.regionwise_engines = regionwise_engines
+
+
+def validate_engine_ref_ids(namespace):
+    """Extracts multiple space-separated identities"""
+    if isinstance(namespace.engine_ref_ids, list):
+        for item in namespace.engine_ref_ids:
+            if not is_valid_resource_id(item):
+                raise InvalidArgumentValueError(f"Invalid engine-ref-ids value: {item}")
+
+
+def validate_engine_ref_ids_and_type(incoming_engine_ref_id_type, engine_ref_ids, exisiting_engine_ref_id_type=None):
+    """Validates combination of engine-ref-id-type and engine-ref-ids"""
+
+    # if engine_ref_id_type is None or SystemAssigned, then no value for engine_ref_ids is expected:
+    engine_ref_id_type = incoming_engine_ref_id_type or exisiting_engine_ref_id_type
+    if engine_ref_id_type != EngineIdentityType.UserAssigned and engine_ref_ids:
+        raise InvalidArgumentValueError(
+            "engine-ref-ids should not be provided when engine-ref-id-type is None or SystemAssigned"
+        )
+
+    # If engine_ref_id_type is UserAssigned, then engine_ref_ids is expected.
+    if incoming_engine_ref_id_type == EngineIdentityType.UserAssigned and engine_ref_ids is None:
+        raise InvalidArgumentValueError(
+            "Atleast one engine-ref-ids should be provided when engine-ref-id-type is UserAssigned"
+        )