[load] API version refresh - 2025-03-01-preview (#8682)

* Add --autostop-engine-users

* Format and organize imports

* Update validators

* Add invalid test case

* Update vendored SDK and bump version

* Fix vendored SDK imports

* Fix trigger

* Mask secrets

Author: Harshan01

src/load/HISTORY.rst

@@ -3,6 +3,11 @@
 Release History
 ===============
 
+2.1.0
+++++++
+* Add option for `--autostop-engine-users` to set maximum users per engine for AutoStop criteria.
+* Update API version to 2025-03-01-preview.
+
 2.0.0
 ++++++
 * Add commands for creating and managing notification rules using CLI. 

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

@@ -109,7 +109,7 @@ def list_notification_rules(
     if test_ids:
         test_ids = ",".join(test_ids)
         logger.info("Filtering notification rules by test ids: %s", test_ids)
-    responses = client.list_notification_rule(test_ids=test_ids)
+    responses = client.list_notification_rules(test_ids=test_ids)
     logger.info("Retrieved notification rules: %s", responses)
     return [response.as_dict() for response in responses]
 

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

@@ -60,6 +60,7 @@ def create_test(
     autostop=None,
     autostop_error_rate=None,
     autostop_error_rate_time_window=None,
+    autostop_maximum_virtual_users_per_engine=None,
     regionwise_engines=None,
     engine_ref_id_type=None,
     engine_ref_ids=None,
@@ -80,7 +81,10 @@ def create_test(
     yaml, yaml_test_body = None, None
     app_components, add_defaults_to_app_components, server_metrics = None, None, None
     autostop_criteria = create_autostop_criteria_from_args(
-        autostop=autostop, error_rate=autostop_error_rate, time_window=autostop_error_rate_time_window)
+        autostop=autostop,
+        error_rate=autostop_error_rate,
+        time_window=autostop_error_rate_time_window,
+        max_vu_per_engine=autostop_maximum_virtual_users_per_engine)
     if load_test_config_file is None:
         test_type = test_type or infer_test_type_from_test_plan(test_plan)
         logger.debug("Inferred test type: %s", test_type)
@@ -198,6 +202,7 @@ def update_test(
     autostop=None,
     autostop_error_rate=None,
     autostop_error_rate_time_window=None,
+    autostop_maximum_virtual_users_per_engine=None,
     regionwise_engines=None,
     engine_ref_id_type=None,
     engine_ref_ids=None,
@@ -215,7 +220,10 @@ def update_test(
     yaml, yaml_test_body = None, None
     app_components, server_metrics, add_defaults_to_app_components = None, None, None
     autostop_criteria = create_autostop_criteria_from_args(
-        autostop=autostop, error_rate=autostop_error_rate, time_window=autostop_error_rate_time_window)
+        autostop=autostop,
+        error_rate=autostop_error_rate,
+        time_window=autostop_error_rate_time_window,
+        max_vu_per_engine=autostop_maximum_virtual_users_per_engine)
     if load_test_config_file is not None:
         yaml = load_yaml(load_test_config_file)
         yaml_test_body = convert_yaml_to_test(cmd, yaml)

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

@@ -34,7 +34,7 @@
         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-error-rate 80.5 --autostop-time-window 120 --autostop-engine-users 1000
         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
     - name: Create a test with a multi-region load configuration using region names in the format accepted by Azure Resource Manager (ARM). Ensure the specified regions are supported by Azure Load Testing. Multi-region load tests are restricted to public endpoints only.

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

@@ -36,6 +36,7 @@ def load_arguments(self, _):
         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)
+        c.argument("autostop_maximum_virtual_users_per_engine", argtypes.autostop_maximum_virtual_users_per_engine)
         c.argument("regionwise_engines", argtypes.regionwise_engines)
         c.argument("engine_ref_id_type", argtypes.engine_ref_id_type)
         c.argument("engine_ref_ids", argtypes.engine_ref_ids)
@@ -62,6 +63,7 @@ def load_arguments(self, _):
         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)
+        c.argument("autostop_maximum_virtual_users_per_engine", argtypes.autostop_maximum_virtual_users_per_engine)
         c.argument("regionwise_engines", argtypes.regionwise_engines)
         c.argument("engine_ref_id_type", argtypes.engine_ref_id_type)
         c.argument("engine_ref_ids", argtypes.engine_ref_ids)

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

@@ -242,6 +242,6 @@ def list_trigger_schedules(
     if test_ids:
         test_ids = ",".join(test_ids)
     logger.info("Schedule trigger states: %s", trigger_states)
-    response_list = client.list_trigger(test_ids=test_ids, states=trigger_states)
+    response_list = client.list_triggers(test_ids=test_ids, states=trigger_states)
     logger.debug("Fetched list of schedule triggers: %s", response_list)
     return [response.as_dict() for response in response_list]

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

@@ -395,6 +395,13 @@
     help="Time window during which the error percentage should be evaluated in seconds.",
 )
 
+autostop_maximum_virtual_users_per_engine = CLIArgumentType(
+    options_list=["--autostop-engine-users"],
+    type=int,
+    validator=validators.validate_autostop_maximum_virtual_users_per_engine,
+    help="Maximum number of users per engine at which the test will stop.",
+)
+
 regionwise_engines = CLIArgumentType(
     options_list=["--regionwise-engines"],
     validator=validators.validate_regionwise_engines,

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

@@ -24,6 +24,7 @@ class LoadTestConfigKeys:
     AUTOSTOP = "autoStop"
     AUTOSTOP_ERROR_RATE = "errorPercentage"
     AUTOSTOP_ERROR_RATE_TIME_WINDOW = "timeWindow"
+    AUTOSTOP_MAX_VU_PER_ENGINE = "maximumVirtualUsersPerEngine"
     FAILURE_CRITERIA = "failureCriteria"
     CLIENT_METRICS_PF = "clientMetrics"
     SERVER_METRICS_PF = "serverMetrics"

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

@@ -274,15 +274,17 @@ 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):
+def create_autostop_criteria_from_args(autostop, error_rate, time_window, max_vu_per_engine):
+    if (autostop is None and error_rate is None and time_window is None and max_vu_per_engine 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
+    if max_vu_per_engine is not None:
+        autostop_criteria["maximumVirtualUsersPerEngine"] = max_vu_per_engine
     return autostop_criteria
 
 
@@ -591,10 +593,16 @@ def create_or_update_test_with_config(
     ):
         new_body["autoStopCriteria"]["errorRateTimeWindowInSeconds"] = \
             body["autoStopCriteria"]["errorRateTimeWindowInSeconds"]
+    if (
+        new_body["autoStopCriteria"].get("maximumVirtualUsersPerEngine") is None
+        and body.get("autoStopCriteria", {}).get("maximumVirtualUsersPerEngine") is not None
+    ):
+        new_body["autoStopCriteria"]["maximumVirtualUsersPerEngine"] = \
+            body["autoStopCriteria"]["maximumVirtualUsersPerEngine"]
 
     if (new_body["autoStopCriteria"].get("autoStopDisabled") is True):
         logger.warning(
-            "Auto stop is disabled. Error rate and time window will be ignored. "
+            "Auto stop is disabled. Error rate, time window and engine users will be ignored. "
             "This can lead to incoming charges for an incorrectly configured test."
         )
 
@@ -749,9 +757,15 @@ def create_or_update_test_without_config(
     ):
         new_body["autoStopCriteria"]["errorRateTimeWindowInSeconds"] = \
             body["autoStopCriteria"]["errorRateTimeWindowInSeconds"]
+    if (
+        new_body["autoStopCriteria"].get("maximumVirtualUsersPerEngine") is None
+        and body.get("autoStopCriteria", {}).get("maximumVirtualUsersPerEngine") is not None
+    ):
+        new_body["autoStopCriteria"]["maximumVirtualUsersPerEngine"] = \
+            body["autoStopCriteria"]["maximumVirtualUsersPerEngine"]
     if (new_body["autoStopCriteria"].get("autoStopDisabled") is True):
         logger.warning(
-            "Auto stop is disabled. Error rate and time window will be ignored. "
+            "Auto stop is disabled. Error rate, time window and engine users will be ignored. "
             "This can lead to incoming charges for an incorrectly configured test."
         )
     new_body["baselineTestRunId"] = baseline_test_run_id if baseline_test_run_id else body.get("baselineTestRunId")

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

@@ -28,15 +28,18 @@ def yaml_parse_autostop_criteria(data):
         }
     error_rate = data[LoadTestConfigKeys.AUTOSTOP].get(LoadTestConfigKeys.AUTOSTOP_ERROR_RATE)
     time_window = data[LoadTestConfigKeys.AUTOSTOP].get(LoadTestConfigKeys.AUTOSTOP_ERROR_RATE_TIME_WINDOW)
+    max_vu_per_engine = data[LoadTestConfigKeys.AUTOSTOP].get(LoadTestConfigKeys.AUTOSTOP_MAX_VU_PER_ENGINE)
     # pylint: disable-next=protected-access
-    validators._validate_autostop_criteria_configfile(error_rate, time_window)
+    validators._validate_autostop_criteria_configfile(error_rate, time_window, max_vu_per_engine)
     autostop_criteria = {
         "autoStopDisabled": False,
     }
     if error_rate is not None:
         autostop_criteria["errorRate"] = error_rate
     if time_window is not None:
         autostop_criteria["errorRateTimeWindowInSeconds"] = time_window
+    if max_vu_per_engine is not None:
+        autostop_criteria["maximumVirtualUsersPerEngine"] = max_vu_per_engine
     return autostop_criteria