[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 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,10 @@
     - 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 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-error-rate-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
 """
 
 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

@@ -12,6 +12,7 @@
     quotes,
     resource_group_name_type,
 )
+from decimal import Decimal
 from knack.arguments import CLIArgumentType
 
 quote_text = f"Use {quotes} to clear existing {{}}."
@@ -341,3 +342,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. The default value is enable.",
+)
+
+autostop_error_rate = CLIArgumentType(
+    options_list=["--autostop-error-rate"],
+    type=Decimal,
+    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

@@ -18,6 +18,7 @@
     CLIInternalError,
 )
 from azure.mgmt.core.tools import is_valid_resource_id, parse_resource_id
+from decimal import Decimal
 from knack.log import get_logger
 
 from .models import IdentityType, AllowedFileTypes
@@ -257,6 +258,16 @@ 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
+    return {
+        "autoStopDisabled": not autostop if autostop is not None else False,
+        "errorRate": error_rate,
+        "errorRateTimeWindowInSeconds": time_window,
+    }
+
+
 def load_yaml(file_path):
     logger.debug("Loading yaml file: %s", file_path)
     try:
@@ -347,6 +358,24 @@ 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)):
+            new_body["autoStopCriteria"] = {
+                "autoStopDisabled": True,
+            }
+        else:
+            error_rate = data["autoStop"].get("errorPercentage", Decimal(90.0))
+            time_window = data["autoStop"].get("timeWindow", 60)
+            new_body["autoStopCriteria"] = {
+                "autoStopDisabled": False,
+                "errorRate": error_rate,
+                "errorRateTimeWindowInSeconds": time_window,
+            }
+    if (new_body.get("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("Converted yaml to test body: %s", new_body)
     return new_body
 
@@ -367,6 +396,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 +503,27 @@ def create_or_update_test_with_config(
         new_body["loadTestConfiguration"]["splitAllCSVs"] = yaml_test_body[
             "loadTestConfiguration"
         ]["splitAllCSVs"]
+
+    logger.debug("autostop_criteria: %s", autostop_criteria)
+    logger.debug("yaml_test_body: %s", yaml_test_body["autoStopCriteria"])
+
+    if autostop_criteria is not None:
+        new_body["autoStopCriteria"] = {
+            "autoStopDisabled": autostop_criteria.get("autoStopDisabled", False),
+            "errorRate": autostop_criteria.get("errorRate", Decimal(90.0)),
+            "errorRateTimeWindowInSeconds": autostop_criteria.get("errorRateTimeWindowInSeconds", 60),
+        }
+    elif yaml_test_body.get("autoStopCriteria") is not None:
+        new_body["autoStopCriteria"] = yaml_test_body["autoStopCriteria"]
+    elif body.get("autoStopCriteria") is not None:
+        new_body["autoStopCriteria"] = body["autoStopCriteria"]
+
+    if (new_body.get("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 +543,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 +610,16 @@ def create_or_update_test_without_config(
         ]["splitAllCSVs"]
     if disable_public_ip is not None:
         new_body["publicIPDisabled"] = disable_public_ip
+
+    if autostop_criteria is not None:
+        new_body["autoStopCriteria"] = {
+            "autoStopDisabled": autostop_criteria.get("autoStopDisabled", False),
+            "errorRate": autostop_criteria.get("errorRate", Decimal(90.0)),
+            "errorRateTimeWindowInSeconds": autostop_criteria.get("errorRateTimeWindowInSeconds", 60),
+        }
+    elif body.get("autoStopCriteria") is not None:
+        new_body["autoStopCriteria"] = body["autoStopCriteria"]
+
     logger.debug("Request body for create or update test: %s", new_body)
     return new_body
 

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

@@ -12,6 +12,7 @@
 from azure.cli.core.azclierror import InvalidArgumentValueError, FileOperationError
 from azure.mgmt.core.tools import is_valid_resource_id
 from knack.log import get_logger
+from decimal import Decimal
 
 from . import utils
 from .models import AllowedFileTypes, AllowedIntervals, AllowedMetricNamespaces
@@ -402,3 +403,42 @@ 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-error-rate-time-window type: {type(namespace.autostop_error_rate_time_window)}"
+        )
+    if namespace.autostop_error_rate_time_window < 1:
+        raise InvalidArgumentValueError(
+            "Autostop error rate time window should be greater than 0"
+        )
+
+
+def validate_autostop_error_rate(namespace):
+    if namespace.autostop_error_rate is None:
+        return
+    if not isinstance(namespace.autostop_error_rate, Decimal):
+        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"
+        )

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

@@ -72,6 +72,11 @@ 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")
+    AUTOSTOP_DISABLED = "disable"
+    AUTOSTOP_ERROR_RATE = 77.5
+    AUTOSTOP_ERROR_RATE_TIME_WINDOW = 90
+
 
 class LoadTestConstants(LoadConstants):
     # Test IDs for load test commands