Adding the CICD improvement changes for the CLI.

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

@@ -21,6 +21,10 @@
     load_yaml,
     upload_file_to_test,
     upload_files_helper,
+    merge_existing_app_components,
+    merge_existing_server_metrics,
+    parse_app_comps_and_server_metrics,
+    is_not_empty_dictionary,
 )
 from azext_load.data_plane.utils.models import (
     AllowedTestTypes,
@@ -48,6 +52,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,
@@ -66,12 +71,14 @@ def create_test(
         body = client.get_test(test_id)
     except ResourceNotFoundError:
         pass
+
     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 +95,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,
@@ -99,6 +107,7 @@ def create_test(
     else:
         yaml = load_yaml(load_test_config_file)
         yaml_test_body = convert_yaml_to_test(cmd, yaml)
+        app_components, add_defaults_to_app_components, server_metrics = parse_app_comps_and_server_metrics(data=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,33 @@ 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 is_not_empty_dictionary(app_components):
+        # 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 is_not_empty_dictionary(server_metrics):
+        # only get and patch the app components if its present in the yaml.
+        server_metrics_existing = None
+        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 +191,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 +214,14 @@ 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)
+        app_components, add_defaults_to_app_components, server_metrics = parse_app_comps_and_server_metrics(data=yaml)
         body = create_or_update_test_with_config(
             test_id,
             body,
@@ -194,6 +233,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 +253,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 +271,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 is_not_empty_dictionary(app_components):
+        # 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 is_not_empty_dictionary(server_metrics):
+        # 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/help.py

@@ -49,6 +49,9 @@
     - 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"
+    - name: Create a test with user assigned Managed Identity reference for accessing the metrics of the configured apps.
+      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 --metrics-reference-id "/subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/sample-rg/providers/microsoft.managedidentity/userassignedidentities/sample-mi"
 """
 
 helps[

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. Provide `null` or `None` to use the system assigned identity of the load test resource."
+        )
         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": ">"
+    }