Adding the CICD improvement changes for the CLI.

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

@@ -21,6 +21,8 @@
     load_yaml,
     upload_file_to_test,
     upload_files_helper,
+    merge_existing_app_components,
+    merge_existing_server_metrics,
 )
 from azext_load.data_plane.utils.models import (
     AllowedTestTypes,
@@ -48,6 +50,7 @@ def create_test(
     secrets=None,
     certificate=None,
     key_vault_reference_identity=None,
+    metrics_reference_identity=None,
     subnet_id=None,
     split_csv=None,
     disable_public_ip=None,
@@ -62,16 +65,21 @@ def create_test(
     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)
     body = None
+    app_components_existing, server_metrics_existing = None, None
     try:
         body = client.get_test(test_id)
     except ResourceNotFoundError:
         pass
+
+    # adding temporarily to avoid the error that the var is not getting used.
+    logger.debug(app_components_existing, server_metrics_existing)
     if body is not None:
         msg = f"Test with given test ID : {test_id} already exist."
         logger.debug(msg)
         raise InvalidArgumentValueError(msg)
     body = {}
     yaml, yaml_test_body = None, None
+    app_components, server_metrics = None, None
     autostop_criteria = create_autostop_criteria_from_args(
         autostop=autostop, error_rate=autostop_error_rate, time_window=autostop_error_rate_time_window)
     if load_test_config_file is None:
@@ -88,6 +96,7 @@ def create_test(
             secrets=secrets,
             certificate=certificate,
             key_vault_reference_identity=key_vault_reference_identity,
+            metrics_reference_identity=metrics_reference_identity,
             subnet_id=subnet_id,
             split_csv=split_csv,
             disable_public_ip=disable_public_ip,
@@ -98,7 +107,7 @@ def create_test(
         )
     else:
         yaml = load_yaml(load_test_config_file)
-        yaml_test_body = convert_yaml_to_test(cmd, yaml)
+        yaml_test_body, app_components, add_defaults_to_app_components, server_metrics = convert_yaml_to_test(cmd, yaml)
         test_type = (
             test_type or
             yaml.get(LoadTestConfigKeys.TEST_TYPE) or
@@ -118,6 +127,7 @@ def create_test(
             secrets=secrets,
             certificate=certificate,
             key_vault_reference_identity=key_vault_reference_identity,
+            metrics_reference_identity=metrics_reference_identity,
             subnet_id=subnet_id,
             split_csv=split_csv,
             disable_public_ip=disable_public_ip,
@@ -136,8 +146,32 @@ def create_test(
     upload_files_helper(
         client, test_id, yaml, test_plan, load_test_config_file, not custom_no_wait, evaluated_test_type
     )
-    response = client.get_test(test_id)
     logger.info("Upload files to test %s has completed", test_id)
+    if app_components is not None and len(app_components) > 0:
+        # only get and patch the app components if its present in the yaml.
+        app_component_response = client.create_or_update_app_components(
+            test_id=test_id, body={"testId": test_id, "components": app_components}
+        )
+        logger.warning(
+            "Added app components for test ID: %s and response is %s", test_id, app_component_response
+        )
+    if server_metrics is not None and len(server_metrics) > 0:
+        # only get and patch the app components if its present in the yaml.
+        try:
+            server_metrics_existing = client.get_server_metrics_config(test_id)
+        except Exception:
+            server_metrics_existing = {"metrics": {}}
+            pass
+        server_metrics_merged = merge_existing_server_metrics(
+            add_defaults_to_app_components, server_metrics, server_metrics_existing.get("metrics", {})
+        )
+        server_metric_response = client.create_or_update_server_metrics_config(
+            test_id=test_id, body={"testId": test_id, "metrics": server_metrics_merged}
+        )
+        logger.warning(
+            "Added server metrics for test ID: %s and response is %s", test_id, server_metric_response
+        )
+    response = client.get_test(test_id)
     logger.info("Test %s has been created successfully", test_id)
     return response.as_dict()
 
@@ -156,6 +190,7 @@ def update_test(
     secrets=None,
     certificate=None,
     key_vault_reference_identity=None,
+    metrics_reference_identity=None,
     subnet_id=None,
     split_csv=None,
     disable_public_ip=None,
@@ -178,11 +213,13 @@ def update_test(
     logger.debug("Retrieved test with test ID: %s and body : %s", test_id, body)
 
     yaml, yaml_test_body = None, None
+    app_components, server_metrics = None, None
     autostop_criteria = create_autostop_criteria_from_args(
         autostop=autostop, error_rate=autostop_error_rate, time_window=autostop_error_rate_time_window)
+    add_defaults_to_app_components = None
     if load_test_config_file is not None:
         yaml = load_yaml(load_test_config_file)
-        yaml_test_body = convert_yaml_to_test(cmd, yaml)
+        yaml_test_body, app_components, add_defaults_to_app_components, server_metrics = convert_yaml_to_test(cmd, yaml)
         body = create_or_update_test_with_config(
             test_id,
             body,
@@ -194,6 +231,7 @@ def update_test(
             secrets=secrets,
             certificate=certificate,
             key_vault_reference_identity=key_vault_reference_identity,
+            metrics_reference_identity=metrics_reference_identity,
             subnet_id=subnet_id,
             split_csv=split_csv,
             disable_public_ip=disable_public_ip,
@@ -213,6 +251,7 @@ def update_test(
             secrets=secrets,
             certificate=certificate,
             key_vault_reference_identity=key_vault_reference_identity,
+            metrics_reference_identity=metrics_reference_identity,
             subnet_id=subnet_id,
             split_csv=split_csv,
             disable_public_ip=disable_public_ip,
@@ -230,6 +269,39 @@ def update_test(
     upload_files_helper(
         client, test_id, yaml, test_plan, load_test_config_file, not custom_no_wait, body.get("kind")
     )
+
+    if app_components is not None and len(app_components) > 0:
+        # only get and patch the app components if its present in the yaml.
+        try:
+            app_components_existing = client.get_app_components(test_id)
+        except Exception:
+            app_components_existing = {"components": {}}
+            pass
+        app_components_merged = merge_existing_app_components(
+            app_components, app_components_existing.get("components", {})
+        )
+        app_component_response = client.create_or_update_app_components(
+            test_id=test_id, body={"testId": test_id, "components": app_components_merged}
+        )
+        logger.warning(
+            "Added app components for test ID: %s and response is %s", test_id, app_component_response
+        )
+    if server_metrics is not None and len(server_metrics) > 0:
+        # only get and patch the app components if its present in the yaml.
+        try:
+            server_metrics_existing = client.get_server_metrics_config(test_id)
+        except Exception:
+            server_metrics_existing = {"metrics": {}}
+            pass
+        server_metrics_merged = merge_existing_server_metrics(
+            add_defaults_to_app_components, server_metrics_existing.get("metrics", {}), server_metrics
+        )
+        server_metric_response = client.create_or_update_server_metrics_config(
+            test_id=test_id, body={"testId": test_id, "metrics": server_metrics_merged}
+        )
+        logger.warning(
+            "Added server metrics for test ID: %s and response is %s", test_id, server_metric_response
+        )
     response = client.get_test(test_id)
     logger.info("Upload files to test %s has completed", test_id)
     logger.info("Test %s has been updated successfully", test_id)

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

@@ -27,6 +27,9 @@ def load_arguments(self, _):
         c.argument(
             "key_vault_reference_identity", argtypes.key_vault_reference_identity
         )
+        c.argument(
+            "metrics_reference_identity", argtypes.metrics_reference_identity, help="The identity that will be used to access the metrics. This will be defaulted to SystemAssigned if not given."
+        )
         c.argument("engine_instances", argtypes.engine_instances)
         c.argument("custom_no_wait", argtypes.custom_no_wait)
         c.argument("disable_public_ip", argtypes.disable_public_ip)
@@ -48,6 +51,9 @@ def load_arguments(self, _):
         c.argument(
             "key_vault_reference_identity", argtypes.key_vault_reference_identity, help="The identity that will be used to access the key vault. Provide `null` or `None` to use the system assigned identity of the load test resource."
         )
+        c.argument(
+            "metrics_reference_identity", argtypes.metrics_reference_identity, help="The identity that will be used to access the metrics. This will be defaulted to SystemAssigned if not given."
+        )
         c.argument("engine_instances", argtypes.engine_instances)
         c.argument("subnet_id", argtypes.subnet_id)
         c.argument("split_csv", argtypes.split_csv)

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

@@ -129,11 +129,19 @@
 )
 
 key_vault_reference_identity = CLIArgumentType(
+    validator=validators.validate_keyvault_identity_ref_id,
     options_list=["--keyvault-reference-id"],
     type=str,
     help="The identity that will be used to access the key vault.",
 )
 
+metrics_reference_identity = CLIArgumentType(
+    validator=validators.validate_metrics_identity_ref_id,
+    options_list=["--metrics-reference-id"],
+    type=str,
+    help="The identity that will be used to get the metrics of the configured apps from server pass-fail criteria.",
+)
+
 split_csv = CLIArgumentType(
     validator=validators.validate_split_csv,
     options_list=["--split-csv"],

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

@@ -24,12 +24,25 @@ class LoadTestConfigKeys:
     AUTOSTOP_ERROR_RATE = "errorPercentage"
     AUTOSTOP_ERROR_RATE_TIME_WINDOW = "timeWindow"
     FAILURE_CRITERIA = "failureCriteria"
+    CLIENT_METRICS_PF = "clientMetrics"
+    SERVER_METRICS_PF = "serverMetrics"
+    METRIC_NAME = "metricName"
+    METRIC_NAME_SERVER_METRICS = "name"
+    METRIC_NAMESPACE_SERVER_METRICS = "namespace"
+    METRIC_NAMESPACE = "metricNamespace"
+    RESOURCEID = "resourceId"
+    AGGREGATION = "aggregation"
+    CONDITION = "condition"
+    APP_COMPONENTS = "appComponents"
+    SERVER_METRICS_APP_COMPONENTS = "metrics"
     REGIONAL_LOADTEST_CONFIG = "regionalLoadTestConfig"
     REGION = "region"
     QUICK_START = "quickStartTest"
     SPLIT_CSV = "splitAllCSVs"
     REFERENCE_IDENTITIES = "referenceIdentities"
     ENGINE = "Engine"
+    METRICS = "Metrics"
+    KEY_VAULT = "KeyVault"
     TYPE = "type"
     KIND = "kind"
     VALUE = "value"
@@ -76,3 +89,11 @@ class LoadTestTrendsKeys:
         AllowedTrendsResponseTimeAggregations.P999.value: "pct999ResTime",
         AllowedTrendsResponseTimeAggregations.P9999.value: "pct9999ResTime",
     }
+
+
+@dataclass
+class LoadTestFailureCriteriaKeys:
+    CONDITION_ENUM_MAP = {
+        "LessThan": "<",
+        "GreaterThan": ">"
+    }

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

@@ -318,6 +318,10 @@ def convert_yaml_to_test(cmd, data):
         new_body["kind"] = data[LoadTestConfigKeys.TEST_TYPE]
     new_body["keyvaultReferenceIdentityType"] = IdentityType.SystemAssigned
     if LoadTestConfigKeys.KEYVAULT_REFERENCE_IDENTITY in data:
+        if not is_valid_resource_id(data[LoadTestConfigKeys.KEYVAULT_REFERENCE_IDENTITY]):
+            raise InvalidArgumentValueError(
+                "Key vault reference identity should be a valid resource id."
+            )
         new_body["keyvaultReferenceIdentityId"] = data[LoadTestConfigKeys.KEYVAULT_REFERENCE_IDENTITY]
         new_body["keyvaultReferenceIdentityType"] = IdentityType.UserAssigned
 
@@ -340,9 +344,10 @@ def convert_yaml_to_test(cmd, 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)
+    utils_yaml_config.update_reference_identities(new_body, data)
+    app_components, add_defaults_to_app_components, server_metrics = utils_yaml_config.parse_app_comps_and_server_metrics(data=data)
     logger.debug("Converted yaml to test body: %s", new_body)
-    return new_body
+    return new_body, app_components, add_defaults_to_app_components, server_metrics
 # pylint: enable=line-too-long
 
 
@@ -360,6 +365,7 @@ def create_or_update_test_with_config(
     secrets=None,
     certificate=None,
     key_vault_reference_identity=None,
+    metrics_reference_identity=None,
     subnet_id=None,
     split_csv=None,
     disable_public_ip=None,
@@ -393,12 +399,23 @@ def create_or_update_test_with_config(
             "keyvaultReferenceIdentityId"
         )
         new_body["keyvaultReferenceIdentityType"] = IdentityType.UserAssigned
-    else:
-        new_body["keyvaultReferenceIdentityType"] = IdentityType.SystemAssigned
     if new_body["keyvaultReferenceIdentityType"] == IdentityType.UserAssigned:
         if new_body["keyvaultReferenceIdentityId"].casefold() in ["null", ""]:
             new_body["keyvaultReferenceIdentityType"] = IdentityType.SystemAssigned
             new_body.pop("keyvaultReferenceIdentityId")
+    new_body["metricsReferenceIdentityType"] = IdentityType.SystemAssigned
+    if metrics_reference_identity is not None:
+        new_body["metricsReferenceIdentityId"] = metrics_reference_identity
+        new_body["metricsReferenceIdentityType"] = IdentityType.UserAssigned
+    elif yaml_test_body.get("metricsReferenceIdentityId") is not None:
+        new_body["metricsReferenceIdentityId"] = yaml_test_body.get(
+            "metricsReferenceIdentityId"
+        )
+        new_body["metricsReferenceIdentityType"] = IdentityType.UserAssigned
+    if new_body["metricsReferenceIdentityType"] == IdentityType.UserAssigned:
+        if new_body["metricsReferenceIdentityId"].casefold() in ["null", ""]:
+            new_body["metricsReferenceIdentityType"] = IdentityType.SystemAssigned
+            new_body.pop("metricsReferenceIdentityId")
     subnet_id = subnet_id or yaml_test_body.get("subnetId")
     if disable_public_ip is not None:
         new_body["publicIPDisabled"] = disable_public_ip
@@ -479,11 +496,18 @@ def create_or_update_test_with_config(
             "passFailMetrics": {
                 key: None
                 for key in existing_pass_fail_Criteria.get("passFailMetrics", {})
+            },
+            "passFailServerMetrics": {
+                key: None
+                for key in existing_pass_fail_Criteria.get("passFailServerMetrics", {})
             }
         }
         new_body["passFailCriteria"]["passFailMetrics"].update(
             yaml_pass_fail_criteria.get("passFailMetrics", {})
         )
+        new_body["passFailCriteria"]["passFailServerMetrics"].update(
+            yaml_pass_fail_criteria.get("passFailServerMetrics", {})
+        )
     if split_csv is not None:
         new_body["loadTestConfiguration"]["splitAllCSVs"] = split_csv
     elif (
@@ -552,6 +576,7 @@ def create_or_update_test_without_config(
     secrets=None,
     certificate=None,
     key_vault_reference_identity=None,
+    metrics_reference_identity=None,
     subnet_id=None,
     split_csv=None,
     disable_public_ip=None,
@@ -585,6 +610,21 @@ def create_or_update_test_without_config(
         if new_body["keyvaultReferenceIdentityId"].casefold() in ["null", ""]:
             new_body["keyvaultReferenceIdentityType"] = IdentityType.SystemAssigned
             new_body.pop("keyvaultReferenceIdentityId")
+    new_body["metricsReferenceIdentityType"] = IdentityType.SystemAssigned
+    if metrics_reference_identity is not None:
+        new_body["metricsReferenceIdentityId"] = metrics_reference_identity
+        new_body["metricsReferenceIdentityType"] = IdentityType.UserAssigned
+    elif body.get("metricsReferenceIdentityId") is not None:
+        new_body["metricsReferenceIdentityId"] = body.get(
+            "metricsReferenceIdentityId"
+        )
+        new_body["metricsReferenceIdentityType"] = body.get(
+            "metricsReferenceIdentityType", IdentityType.UserAssigned
+        )
+    if new_body["metricsReferenceIdentityType"] == IdentityType.UserAssigned:
+        if new_body["metricsReferenceIdentityId"].casefold() in ["null", ""]:
+            new_body["metricsReferenceIdentityType"] = IdentityType.SystemAssigned
+            new_body.pop("metricsReferenceIdentityId")
     subnet_id = subnet_id or body.get("subnetId")
     if subnet_id:
         if subnet_id.casefold() in ["null", ""]:
@@ -931,3 +971,22 @@ def _add_error_and_throughput_trends(trends, test_run):
     throughput = _get_metrics_from_sampler(test_run, "Total", "throughput")
     if throughput is not None:
         trends[LoadTestTrendsKeys.THROUGHPUT] = round(throughput, 2)
+
+
+def merge_existing_app_components(app_components_yaml, existing_app_components):
+    if existing_app_components is None:
+        return app_components_yaml
+    for key in existing_app_components:
+        if key not in app_components_yaml:
+            app_components_yaml[key] = None
+    return app_components_yaml
+
+
+def merge_existing_server_metrics(add_defaults_to_app_copmponents, existing_server_metrics, server_metrics_yaml):
+    if existing_server_metrics is None:
+        return server_metrics_yaml
+    for key in existing_server_metrics:
+        resourceid = (existing_server_metrics.get(key) or {}).get(LoadTestConfigKeys.RESOURCEID, "").lower()
+        if key not in server_metrics_yaml and (add_defaults_to_app_copmponents.get(resourceid) is None or add_defaults_to_app_copmponents.get(resourceid) is False):
+            server_metrics_yaml[key] = None
+    return server_metrics_yaml