[load] support multi-region load test configuration

Author: mbhardwaj-msft

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

@@ -46,6 +46,7 @@ def create_test(
     autostop=None,
     autostop_error_rate=None,
     autostop_error_rate_time_window=None,
+    regionwise_engines=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)
@@ -79,6 +80,7 @@ def create_test(
             split_csv=split_csv,
             disable_public_ip=disable_public_ip,
             autostop_criteria=autostop_criteria,
+            regionwise_engines=regionwise_engines,
         )
     else:
         yaml = load_yaml(load_test_config_file)
@@ -97,7 +99,8 @@ def create_test(
             subnet_id=subnet_id,
             split_csv=split_csv,
             disable_public_ip=disable_public_ip,
-            autostop_criteria=autostop_criteria
+            autostop_criteria=autostop_criteria,
+            regionwise_engines=regionwise_engines,
         )
     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)
@@ -135,6 +138,7 @@ def update_test(
     autostop=None,
     autostop_error_rate=None,
     autostop_error_rate_time_window=None,
+    regionwise_engines=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)
@@ -166,7 +170,8 @@ def update_test(
             subnet_id=subnet_id,
             split_csv=split_csv,
             disable_public_ip=disable_public_ip,
-            autostop_criteria=autostop_criteria
+            autostop_criteria=autostop_criteria,
+            regionwise_engines=regionwise_engines,
         )
     else:
         body = create_or_update_test_without_config(
@@ -182,7 +187,8 @@ def update_test(
             subnet_id=subnet_id,
             split_csv=split_csv,
             disable_public_ip=disable_public_ip,
-            autostop_criteria=autostop_criteria
+            autostop_criteria=autostop_criteria,
+            regionwise_engines=regionwise_engines,
         )
     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/params.py

@@ -32,6 +32,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("regionwise_engines", argtypes.regionwise_engines)
 
     with self.argument_context("load test update") as c:
         c.argument("load_test_config_file", argtypes.load_test_config_file)
@@ -52,6 +53,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("regionwise_engines", argtypes.regionwise_engines)
 
     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

@@ -362,3 +362,10 @@
     validator=validators.validate_autostop_error_rate_time_window,
     help="Time window during which the error percentage should be evaluated in seconds.",
 )
+
+regionwise_engines = CLIArgumentType(
+    options_list=["--regionwise-engines"],
+    validator=validators.validate_regionwise_engines,
+    nargs="+",
+    help="Regionwise engine count in the format of region1=engineCount1 region2=engineCount2 ... The region name should be of format accepted by ARM, and should be a region supported by Azure Load Testing.",
+)

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

@@ -257,6 +257,21 @@ def parse_env(envs):
     return env_dict
 
 
+def parse_regionwise_loadtest_config(regionwise_loadtest_config):
+    logger.debug("Parsing regionwise load test configuration")
+    regional_load_test_config = []
+    for region_load in regionwise_loadtest_config:
+        region_name = region_load.get("region")
+        if region_name is None or not isinstance(region_name, str):
+            raise RequiredArgumentMissingError("Region name is required of type string")
+        engine_instances = region_load.get("engineInstances")
+        if engine_instances is None or not isinstance(engine_instances, int):
+            raise InvalidArgumentValueError("Engine instances is required of type integer")
+        regional_load_test_config.append({"region": region_name.lower(), "engineInstances": engine_instances})
+    logger.debug("Successfully parsed regionwise load test configuration: %s", regional_load_test_config)
+    return regional_load_test_config
+
+
 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
@@ -304,9 +319,12 @@ def convert_yaml_to_test(data):
         new_body["subnetId"] = data["subnetId"]
 
     new_body["loadTestConfiguration"] = {}
-    new_body["loadTestConfiguration"]["engineInstances"] = data.get(
-        "engineInstances", 1
-    )
+    new_body["loadTestConfiguration"]["engineInstances"] = data.get("engineInstances")
+    if data.get("regionalLoadTestConfig") is not None:
+        new_body["loadTestConfiguration"]["regionalLoadTestConfig"] = parse_regionwise_loadtest_config(
+            data.get("regionalLoadTestConfig")
+        )
+
     if data.get("certificates"):
         new_body["certificate"] = parse_cert(data.get("certificates"))
     if data.get("secrets"):
@@ -399,6 +417,7 @@ def create_or_update_test_with_config(
     split_csv=None,
     disable_public_ip=None,
     autostop_criteria=None,
+    regionwise_engines=None,
 ):
     logger.info(
         "Creating a request body for create or update test using config and parameters."
@@ -480,7 +499,25 @@ def create_or_update_test_with_config(
             "loadTestConfiguration"
         ]["engineInstances"]
     else:
-        new_body["loadTestConfiguration"]["engineInstances"] = 1
+        new_body["loadTestConfiguration"]["engineInstances"] = body.get(
+            "loadTestConfiguration", {}
+        ).get("engineInstances", 1)
+    if regionwise_engines:
+        new_body["loadTestConfiguration"]["regionalLoadTestConfig"] = regionwise_engines
+    elif (
+        yaml_test_body.get("loadTestConfiguration", {}).get("regionalLoadTestConfig")
+        is not None
+    ):
+        new_body["loadTestConfiguration"]["regionalLoadTestConfig"] = yaml_test_body[
+            "loadTestConfiguration"
+        ]["regionalLoadTestConfig"]
+    else:
+        new_body["loadTestConfiguration"]["regionalLoadTestConfig"] = body.get(
+            "loadTestConfiguration", {}
+        ).get("regionalLoadTestConfig")
+    validate_engine_data_with_regionwiseload_data(
+        new_body["loadTestConfiguration"]["engineInstances"],
+        new_body["loadTestConfiguration"]["regionalLoadTestConfig"])
     # quick test is not supported in CLI
     new_body["loadTestConfiguration"]["quickStartTest"] = False
 
@@ -554,6 +591,7 @@ def create_or_update_test_without_config(
     split_csv=None,
     disable_public_ip=None,
     autostop_criteria=None,
+    regionwise_engines=None,
 ):
     logger.info(
         "Creating a request body for test using parameters and old test body (in case of update)."
@@ -610,6 +648,15 @@ def create_or_update_test_without_config(
         new_body["loadTestConfiguration"]["engineInstances"] = body.get(
             "loadTestConfiguration", {}
         ).get("engineInstances", 1)
+    if regionwise_engines:
+        new_body["loadTestConfiguration"]["regionalLoadTestConfig"] = regionwise_engines
+    else:
+        new_body["loadTestConfiguration"]["regionalLoadTestConfig"] = body.get(
+            "loadTestConfiguration", {}
+        ).get("regionalLoadTestConfig")
+    validate_engine_data_with_regionwiseload_data(
+        new_body["loadTestConfiguration"]["engineInstances"],
+        new_body["loadTestConfiguration"]["regionalLoadTestConfig"])
     # quick test is not supported in CLI
     new_body["loadTestConfiguration"]["quickStartTest"] = False
     if split_csv is not None:
@@ -804,6 +851,19 @@ def upload_files_helper(
         load_test_config_file=load_test_config_file, existing_test_files=files, wait=wait)
 
 
+def validate_engine_data_with_regionwiseload_data(engine_instances, regionwise_engines):
+    if regionwise_engines is None:
+        return
+    total_engines = 0
+    for region in regionwise_engines:
+        total_engines += region["engineInstances"]
+    if total_engines != engine_instances:
+        raise InvalidArgumentValueError(
+            f"Sum of engine instances in regionwise load test configuration ({total_engines}) "
+            f"should be equal to total engine instances ({engine_instances})"
+        )
+
+
 def validate_failure_criteria(failure_criteria):
     parts = failure_criteria.split("(")
     if len(parts) != 2:

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

@@ -464,3 +464,32 @@ def _validate_autostop_criteria_configfile(error_rate, time_window):
         raise InvalidArgumentValueError(
             "Invalid value for timeWindow. Value should be an integer greater than or equal to 0"
         )
+
+
+def validate_regionwise_engines(namespace):
+    if namespace.regionwise_engines is None:
+        return
+    if not isinstance(namespace.regionwise_engines, list):
+        raise InvalidArgumentValueError(
+            f"Invalid regionwise-engines type: {type(namespace.regionwise_engines)}. \
+                Expected list in the format of region1=engineCount1 region2=engineCount2"
+        )
+    regionwise_engines = []
+    for item in namespace.regionwise_engines:
+        if not isinstance(item, str) or "=" not in item:
+            raise InvalidArgumentValueError(
+                f"Invalid regionwise-engines item type: {type(item)}. Expected region=engineCount"
+            )
+        key, value = item.split("=", 1)
+        if not key or not value:
+            raise InvalidArgumentValueError(
+                f"Invalid regionwise-engines item: {item}. Region or engine count cannot be empty"
+            )
+        try:
+            value = int(value.strip())
+        except ValueError:
+            raise InvalidArgumentValueError(
+                f"Invalid regionwise-engines item value: {value}. Expected integer"
+            )
+        regionwise_engines.append({"region": key.strip().lower(), "engineInstances": value})
+    namespace.regionwise_engines = regionwise_engines

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

@@ -33,6 +33,27 @@ class LoadConstants:
     INVALID_ZIP_ARTIFACT_WITH_SUBDIR_NAME = "sample-ZIP-artifact-subdir.zip"
     INVALID_ZIP_ARTIFACT_WITH_SUBDIR_FILE = os.path.join(TEST_RESOURCES_DIR, r"sample-ZIP-artifact-subdir.zip")
 
+    # Constants for Regional Load Config Unit Tests
+    REGIONAL_LOAD_CONFIG_FILE = os.path.join(TEST_RESOURCES_DIR, r"config-regionwise-engines.yaml")
+    REGIONAL_LOAD_CONFIG_FILE_COUNT_MISMATCH = os.path.join(TEST_RESOURCES_DIR, r"config-regionwise-engines-count-mismatch.yaml")
+    REGIONAL_LOAD_CONFIG_FILE_INVALID_REGION = os.path.join(TEST_RESOURCES_DIR, r"config-regionwise-engines-invalid-region.yaml")
+    REGIONAL_LOAD_CONFIG_FILE_INVALID_TYPE_FLOAT = os.path.join(TEST_RESOURCES_DIR, r"config-regionwise-engines-invalid-type-float.yaml")
+    REGIONAL_LOAD_CONFIG_FILE_INVALID_TYPE_STRING = os.path.join(TEST_RESOURCES_DIR, r"config-regionwise-engines-invalid-type-string.yaml")
+    REGIONAL_LOAD_CONFIG_FILE_NO_PARENT_REGION = os.path.join(TEST_RESOURCES_DIR, r"config-regionwise-engines-no-parent-region.yaml")
+    REGIONAL_LOAD_CONFIG_FILE_NO_TOTAL = os.path.join(TEST_RESOURCES_DIR, r"config-regionwise-engines-no-total.yaml")
+    ENGINE_INSTANCES = 5
+    REGIONWISE_ENGINES = "germanywestcentral=2 eastus=3"
+    REGIONWISE_ENGINES_1 = "eastus=1"
+    REGIONWISE_ENGINES_2 = '"southcentralus = 4" "eastus = 1"'
+    REGIONWISE_ENGINES_3 = '"southcentralus = 2" "southcentralus = 2" "eastus = 1"'
+    REGIONWISE_ENGINES_INVALID_REGION = "invalidregion=2 eastus=3"
+    REGIONWISE_ENGINES_INVALID_TYPE_FLOAT = "germanywestcentral=2 eastus=3.5"
+    REGIONWISE_ENGINES_INVALID_TYPE_STRING = "germanywestcentral=2 eastus=three"
+    REGIONWISE_ENGINES_NO_PARENT_REGION = "germanywestcentral=2 uksouth=3"
+    REGIONWISE_ENGINES_INVALID_FORMAT_1 = {"germanywestcentral": 2, "eastus": 3}
+    REGIONWISE_ENGINES_INVALID_FORMAT_2 = "germanywestcentral=2 eastus:3"
+    REGIONWISE_ENGINES_INVALID_FORMAT_3 = "=2 eastus=3"
+
     ENV_VAR_DURATION_NAME = "duration_in_sec"
     ENV_VAR_DURATION_SHORT = "1"
     ENV_VAR_DURATION_LONG = "120"
@@ -98,6 +119,7 @@ class LoadTestConstants(LoadConstants):
     APP_COMPONENT_TEST_ID = "app-component-test-case"
     SERVER_METRIC_TEST_ID = "server-metric-test-case"
     FILE_TEST_ID = "file-test-case"
+    REGIONAL_LOAD_CONFIG_TEST_ID = "regional-load-config-test-case"
 
     INVALID_UPDATE_TEST_ID = "invalid-update-test-case"
     INVALID_PF_TEST_ID = "invalid-pf-test-case"