Adding the CICD improvement changes for the CLI.

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

@@ -48,6 +48,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,
@@ -72,6 +73,7 @@ def create_test(
         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 +90,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 +101,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, server_metrics = convert_yaml_to_test(cmd, yaml)
         test_type = (
             test_type or
             yaml.get(LoadTestConfigKeys.TEST_TYPE) or
@@ -118,6 +121,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 +140,22 @@ 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:
+        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 obj is %s", test_id, app_component_response
+        )
+    if server_metrics is not None and len(server_metrics) > 0:
+        server_metric_response = client.create_or_update_server_metrics_config(
+            test_id=test_id, body={"testId": test_id, "metrics": server_metrics}
+        )
+        logger.warning(
+            "Added server metrics for test ID: %s and response obj 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 +174,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 +197,12 @@ 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)
     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, server_metrics = convert_yaml_to_test(cmd, yaml)
         body = create_or_update_test_with_config(
             test_id,
             body,
@@ -194,6 +214,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 +234,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 +252,21 @@ 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:
+        app_component_response = client.create_or_update_app_components(
+            test_id=test_id, body={"testId": test_id, "components": app_components}
+        )
+        logger.info(
+            "Updated app components for test ID: %s and response obj is %s", test_id, app_component_response
+        )
+    server_metric_response = None
+    if server_metrics is not None and len(server_metrics) > 0:
+        server_metric_response = client.create_or_update_server_metrics_config(
+            test_id=test_id, body={"testId": test_id, "metrics": server_metrics}
+        )
+        logger.info(
+            "Updated server metrics for test ID: %s and response obj 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, 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, 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,
@@ -399,6 +405,21 @@ def create_or_update_test_with_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 yaml_test_body.get("metricsReferenceIdentityId") is not None:
+        new_body["metricsReferenceIdentityId"] = yaml_test_body.get(
+            "metricsReferenceIdentityId"
+        )
+        new_body["metricsReferenceIdentityType"] = IdentityType.UserAssigned
+    else:
+        new_body["metricsReferenceIdentityType"] = IdentityType.SystemAssigned
+    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 +500,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 +580,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 +614,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", ""]: