[az][load] support for autostop criteria

src/load/HISTORY.rst

@@ -2,6 +2,10 @@
 
 Release History
 ===============
+1.3.0
+++++++
+* Add support for autostop criteria. Autostop error rate and time window in seconds can be set using `--autostop-error-rate` and `--autostop-time-window` arguments in 'az load test create' and 'az load test update' commands. Autostop can be disabled by using `--autostop disable` in 'az load test create' and 'az load test update' commands. Autostop criteria set in YAML config file will now also be honoured.
+
 1.2.0
 ++++++
 * Added support for disable public IP in test creation and update. This can be done by using --disable-public-ip argument in 'az load test create' and 'az load test update' commands.

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

@@ -16,6 +16,7 @@
     load_yaml,
     upload_file_to_test,
     upload_files_helper,
+    create_autostop_criteria_from_args,
 )
 from azure.cli.core.azclierror import InvalidArgumentValueError
 from azure.core.exceptions import ResourceNotFoundError
@@ -42,6 +43,9 @@ def create_test(
     split_csv=None,
     disable_public_ip=None,
     custom_no_wait=False,
+    autostop=None,
+    autostop_error_rate=None,
+    autostop_error_rate_time_window=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)
@@ -58,6 +62,8 @@ def create_test(
     yaml, yaml_test_body = None, None
     if split_csv is None:
         split_csv = False
+    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:
         body = create_or_update_test_without_config(
             test_id,
@@ -72,6 +78,7 @@ def create_test(
             subnet_id=subnet_id,
             split_csv=split_csv,
             disable_public_ip=disable_public_ip,
+            autostop_criteria=autostop_criteria,
         )
     else:
         yaml = load_yaml(load_test_config_file)
@@ -90,6 +97,7 @@ def create_test(
             subnet_id=subnet_id,
             split_csv=split_csv,
             disable_public_ip=disable_public_ip,
+            autostop_criteria=autostop_criteria
         )
     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)
@@ -124,6 +132,9 @@ def update_test(
     split_csv=None,
     disable_public_ip=None,
     custom_no_wait=False,
+    autostop=None,
+    autostop_error_rate=None,
+    autostop_error_rate_time_window=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)
@@ -136,6 +147,8 @@ def update_test(
     logger.debug("Retrieved test with test ID: %s and body : %s", test_id, body)
 
     yaml, yaml_test_body = 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(yaml)
@@ -153,6 +166,7 @@ def update_test(
             subnet_id=subnet_id,
             split_csv=split_csv,
             disable_public_ip=disable_public_ip,
+            autostop_criteria=autostop_criteria
         )
     else:
         body = create_or_update_test_without_config(
@@ -168,6 +182,7 @@ def update_test(
             subnet_id=subnet_id,
             split_csv=split_csv,
             disable_public_ip=disable_public_ip,
+            autostop_criteria=autostop_criteria
         )
     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

@@ -32,6 +32,11 @@
     - name: Create a test for a private endpoint in a Virtual Network with split CSV option enabled.
       text: |
         az load test create --test-id sample-test-id --load-test-resource sample-alt-resource --resource-group sample-rg --display-name "Sample Name" --subnet-id "/subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/sample-rg/providers/Microsoft.Network/virtualNetworks/SampleVMVNET/subnets/SampleVMSubnet" --split-csv true
+    - name: Create a test with custom defined autostop criteria or enable / disable autostop for a test.
+      text: |
+        az load test create --test-id sample-test-id --load-test-resource sample-alt-resource --resource-group sample-rg --display-name "Sample Name" --autostop-error-rate 80.5 --autostop-time-window 120
+        az load test create --test-id sample-test-id --load-test-resource sample-alt-resource --resource-group sample-rg --display-name "Sample Name" --autostop disable
+        az load test create --test-id sample-test-id --load-test-resource sample-alt-resource --resource-group sample-rg --display-name "Sample Name" --autostop enable
 """
 
 helps[

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

@@ -29,6 +29,9 @@ def load_arguments(self, _):
         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)
+        c.argument("autostop", argtypes.autostop)
+        c.argument("autostop_error_rate", argtypes.autostop_error_rate)
+        c.argument("autostop_error_rate_time_window", argtypes.autostop_error_rate_time_window)
 
     with self.argument_context("load test update") as c:
         c.argument("load_test_config_file", argtypes.load_test_config_file)
@@ -46,6 +49,9 @@ def load_arguments(self, _):
         c.argument("split_csv", argtypes.split_csv)
         c.argument("custom_no_wait", argtypes.custom_no_wait)
         c.argument("disable_public_ip", argtypes.disable_public_ip)
+        c.argument("autostop", argtypes.autostop)
+        c.argument("autostop_error_rate", argtypes.autostop_error_rate)
+        c.argument("autostop_error_rate_time_window", argtypes.autostop_error_rate_time_window)
 
     with self.argument_context("load test download-files") as c:
         c.argument("path", argtypes.dir_path)

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

@@ -341,3 +341,24 @@
         "Example: `--dimension-filters key1=value1 key2=*`, `--dimension-filters *`"
     ),
 )
+
+autostop = CLIArgumentType(
+    validator=validators.validate_autostop_enable_disable,
+    options_list=["--autostop"],
+    type=str,
+    help="Whether auto-stop should be enabled or disabled. Allowed values are enable/disable.",
+)
+
+autostop_error_rate = CLIArgumentType(
+    options_list=["--autostop-error-rate"],
+    type=float,
+    validator=validators.validate_autostop_error_rate,
+    help="Threshold percentage of errors on which test run should be automatically stopped. Allowed values are in range of [0.0,100.0]",
+)
+
+autostop_error_rate_time_window = CLIArgumentType(
+    options_list=["--autostop-time-window"],
+    type=int,
+    validator=validators.validate_autostop_error_rate_time_window,
+    help="Time window during which the error percentage should be evaluated in seconds.",
+)

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

@@ -257,6 +257,18 @@ def parse_env(envs):
     return env_dict
 
 
+def create_autostop_criteria_from_args(autostop, error_rate, time_window):
+    if (autostop is None and error_rate is None and time_window is None):
+        return None
+    autostop_criteria = {}
+    autostop_criteria["autoStopDisabled"] = not autostop if autostop is not None else False
+    if error_rate is not None:
+        autostop_criteria["errorRate"] = error_rate
+    if time_window is not None:
+        autostop_criteria["errorRateTimeWindowInSeconds"] = time_window
+    return autostop_criteria
+
+
 def load_yaml(file_path):
     logger.debug("Loading yaml file: %s", file_path)
     try:
@@ -347,6 +359,25 @@ def convert_yaml_to_test(data):
                 new_body["passFailCriteria"]["passFailMetrics"][metric_id][
                     "requestName"
                 ] = name
+    if data.get("autoStop") is not None:
+        if (isinstance(data["autoStop"], str)):
+            # pylint: disable-next=protected-access
+            validators._validate_autostop_disable_configfile(data["autoStop"])
+            new_body["autoStopCriteria"] = {
+                "autoStopDisabled": True,
+            }
+        else:
+            error_rate = data["autoStop"].get("errorPercentage")
+            time_window = data["autoStop"].get("timeWindow")
+            # pylint: disable-next=protected-access
+            validators._validate_autostop_criteria_configfile(error_rate, time_window)
+            new_body["autoStopCriteria"] = {
+                "autoStopDisabled": False,
+            }
+            if error_rate is not None:
+                new_body["autoStopCriteria"]["errorRate"] = error_rate
+            if time_window is not None:
+                new_body["autoStopCriteria"]["errorRateTimeWindowInSeconds"] = time_window
     logger.debug("Converted yaml to test body: %s", new_body)
     return new_body
 
@@ -367,6 +398,7 @@ def create_or_update_test_with_config(
     subnet_id=None,
     split_csv=None,
     disable_public_ip=None,
+    autostop_criteria=None,
 ):
     logger.info(
         "Creating a request body for create or update test using config and parameters."
@@ -473,6 +505,35 @@ def create_or_update_test_with_config(
         new_body["loadTestConfiguration"]["splitAllCSVs"] = yaml_test_body[
             "loadTestConfiguration"
         ]["splitAllCSVs"]
+
+    new_body["autoStopCriteria"] = {}
+    if autostop_criteria is not None:
+        new_body["autoStopCriteria"] = autostop_criteria
+    elif yaml_test_body.get("autoStopCriteria") is not None:
+        new_body["autoStopCriteria"] = yaml_test_body["autoStopCriteria"]
+    if (
+        new_body["autoStopCriteria"].get("autoStopDisabled") is None
+        and body.get("autoStopCriteria", {}).get("autoStopDisabled") is not None
+    ):
+        new_body["autoStopCriteria"]["autoStopDisabled"] = body["autoStopCriteria"]["autoStopDisabled"]
+    if (
+        new_body["autoStopCriteria"].get("errorRate") is None
+        and body.get("autoStopCriteria", {}).get("errorRate") is not None
+    ):
+        new_body["autoStopCriteria"]["errorRate"] = body["autoStopCriteria"]["errorRate"]
+    if (
+        new_body["autoStopCriteria"].get("errorRateTimeWindowInSeconds") is None
+        and body.get("autoStopCriteria", {}).get("errorRateTimeWindowInSeconds") is not None
+    ):
+        new_body["autoStopCriteria"]["errorRateTimeWindowInSeconds"] = \
+            body["autoStopCriteria"]["errorRateTimeWindowInSeconds"]
+
+    if (new_body["autoStopCriteria"].get("autoStopDisabled") is True):
+        logger.warning(
+            "Auto stop is disabled. Error rate and time window will be ignored. "
+            "This can lead to incoming charges for an incorrectly configured test."
+        )
+
     logger.debug("Request body for create or update test: %s", new_body)
     return new_body
 
@@ -492,6 +553,7 @@ def create_or_update_test_without_config(
     subnet_id=None,
     split_csv=None,
     disable_public_ip=None,
+    autostop_criteria=None,
 ):
     logger.info(
         "Creating a request body for test using parameters and old test body (in case of update)."
@@ -558,6 +620,32 @@ def create_or_update_test_without_config(
         ]["splitAllCSVs"]
     if disable_public_ip is not None:
         new_body["publicIPDisabled"] = disable_public_ip
+
+    new_body["autoStopCriteria"] = {}
+    if autostop_criteria is not None:
+        new_body["autoStopCriteria"] = autostop_criteria
+    if (
+        new_body["autoStopCriteria"].get("autoStopDisabled") is None
+        and body.get("autoStopCriteria", {}).get("autoStopDisabled") is not None
+    ):
+        new_body["autoStopCriteria"]["autoStopDisabled"] = body["autoStopCriteria"]["autoStopDisabled"]
+    if (
+        new_body["autoStopCriteria"].get("errorRate") is None
+        and body.get("autoStopCriteria", {}).get("errorRate") is not None
+    ):
+        new_body["autoStopCriteria"]["errorRate"] = body["autoStopCriteria"]["errorRate"]
+    if (
+        new_body["autoStopCriteria"].get("errorRateTimeWindowInSeconds") is None
+        and body.get("autoStopCriteria", {}).get("errorRateTimeWindowInSeconds") is not None
+    ):
+        new_body["autoStopCriteria"]["errorRateTimeWindowInSeconds"] = \
+            body["autoStopCriteria"]["errorRateTimeWindowInSeconds"]
+    if (new_body["autoStopCriteria"].get("autoStopDisabled") is True):
+        logger.warning(
+            "Auto stop is disabled. Error rate and time window will be ignored. "
+            "This can lead to incoming charges for an incorrectly configured test."
+        )
+
     logger.debug("Request body for create or update test: %s", new_body)
     return new_body
 

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

@@ -402,3 +402,65 @@ def validate_disable_public_ip(namespace):
         namespace.disable_public_ip = True
     else:
         namespace.disable_public_ip = False
+
+
+def validate_autostop_enable_disable(namespace):
+    if namespace.autostop is None:
+        return
+    if not isinstance(namespace.autostop, str) or namespace.autostop.casefold() not in ["enable", "disable"]:
+        raise InvalidArgumentValueError(
+            f"Invalid autostop type: {type(namespace.autostop)}. Allowed values: enable, disable"
+        )
+    if namespace.autostop.casefold() not in ["disable"]:
+        namespace.autostop = True
+    else:
+        namespace.autostop = False
+
+
+def validate_autostop_error_rate_time_window(namespace):
+    if namespace.autostop_error_rate_time_window is None:
+        return
+    if not isinstance(namespace.autostop_error_rate_time_window, int):
+        raise InvalidArgumentValueError(
+            f"Invalid autostop-time-window type: {type(namespace.autostop_error_rate_time_window)}"
+        )
+    if namespace.autostop_error_rate_time_window < 0:
+        raise InvalidArgumentValueError(
+            "Autostop error rate time window should be greater than or equal to 0"
+        )
+
+
+def validate_autostop_error_rate(namespace):
+    if namespace.autostop_error_rate is None:
+        return
+    if not isinstance(namespace.autostop_error_rate, float):
+        raise InvalidArgumentValueError(
+            f"Invalid autostop-error-rate type: {type(namespace.autostop_error_rate)}"
+        )
+    if namespace.autostop_error_rate < 0.0 or namespace.autostop_error_rate > 100.0:
+        raise InvalidArgumentValueError(
+            "Autostop error rate should be in range of [0.0,100.0]"
+        )
+
+
+def _validate_autostop_disable_configfile(autostop):
+    if autostop.casefold() not in ["disable"]:
+        raise InvalidArgumentValueError(
+            "Invalid value for autoStop. Valid values are 'disable' or an object with errorPercentage and timeWindow"
+        )
+
+
+def _validate_autostop_criteria_configfile(error_rate, time_window):
+    if error_rate is not None:
+        if isinstance(error_rate, float) and (error_rate < 0.0 or error_rate > 100.0):
+            raise InvalidArgumentValueError(
+                "Invalid value for errorPercentage. Value should be a number between 0.0 and 100.0"
+            )
+        if isinstance(error_rate, int) and (error_rate < 0 or error_rate > 100):
+            raise InvalidArgumentValueError(
+                "Invalid value for errorPercentage. Value should be a number between 0.0 and 100.0"
+            )
+    if time_window is not None and (not isinstance(time_window, int) or time_window < 0):
+        raise InvalidArgumentValueError(
+            "Invalid value for timeWindow. Value should be an integer greater than or equal to 0"
+        )

src/load/azext_load/tests/latest/constants.py

@@ -72,6 +72,17 @@ class LoadConstants:
 
     INVALID_SERVER_METRIC_ID = r"/subscriptions/invalid/resource/id"
 
+    LOAD_TEST_CONFIG_FILE_WITH_AUTOSTOP = os.path.join(TEST_RESOURCES_DIR, r"config-autostop-criteria.yaml")
+    LOAD_TEST_CONFIG_FILE_WITH_AUTOSTOP_ERROR_RATE = os.path.join(TEST_RESOURCES_DIR, r"config-autostop-criteria-error-rate.yaml")
+    LOAD_TEST_CONFIG_FILE_WITH_AUTOSTOP_TIME_WINDOW = os.path.join(TEST_RESOURCES_DIR, r"config-autostop-criteria-time-window.yaml")
+    LOAD_TEST_CONFIG_FILE_WITH_INVALID_AUTOSTOP_ERROR_RATE = os.path.join(TEST_RESOURCES_DIR, r"config-invalid-autostop-criteria-error-rate.yaml")
+    LOAD_TEST_CONFIG_FILE_WITH_INVALID_AUTOSTOP_TIME_WINDOW = os.path.join(TEST_RESOURCES_DIR, r"config-invalid-autostop-criteria-time-window.yaml")
+    LOAD_TEST_CONFIG_FILE_WITH_INVALID_AUTOSTOP = os.path.join(TEST_RESOURCES_DIR, r"config-invalid-autostop-criteria-random-string.yaml")
+    AUTOSTOP_DISABLED = "disable"
+    AUTOSTOP_ERROR_RATE = 77.5
+    AUTOSTOP_ERROR_RATE_INTEGER = 75
+    AUTOSTOP_ERROR_RATE_TIME_WINDOW = 90
+
 
 class LoadTestConstants(LoadConstants):
     # Test IDs for load test commands