[load][CCOA] support multi-region load test configuration | bug fix | refactor 1.0

src/load/HISTORY.rst

@@ -2,6 +2,11 @@
 
 Release History
 ===============
+1.4.0
+++++++
+* Add support for multi-region load test configuration. Multi-region load test configuration can be set using `--regionwise-engines` argument in 'az load test create' and 'az load test update' commands. Multi-region load test configuration set in YAML config file under key `regionalLoadTestConfig` will also be honoured.
+* Bug fix for `engineInstances` being reset to 1 and not getting backfilled using test's existing configuration when engine instances are not explicitly specified either in YAML config file or CLI argument.
+
 1.3.1
 ++++++
 * Bug fix for `splitAllCSVs` not being honoured from config file due to CLI argument being set as false by default leading to configuration not being selected from the config file.

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)
@@ -77,10 +78,11 @@ 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)
-        yaml_test_body = convert_yaml_to_test(yaml)
+        yaml_test_body = convert_yaml_to_test(cmd, yaml)
         body = create_or_update_test_with_config(
             test_id,
             body,
@@ -95,7 +97,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)
@@ -133,6 +136,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)
@@ -149,7 +153,7 @@ def update_test(
         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)
+        yaml_test_body = convert_yaml_to_test(cmd, yaml)
         body = create_or_update_test_with_config(
             test_id,
             body,
@@ -164,7 +168,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(
@@ -180,7 +185,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/help.py

@@ -37,6 +37,9 @@
         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
+    - 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.
+      text: |
+        az load test create --test-id sample-test-id --load-test-resource sample-alt-resource --resource-group sample-rg --engine-instances 3 --regionwise-engines eastus=1 westus2=1 germanywestcentral=1 --test-plan sample-jmx.jmx
 """
 
 helps[
@@ -82,6 +85,12 @@
     - name: Update the Key Vault reference identity to system assigned Managed Identity.
       text: |
         az load test update --load-test-resource sample-alt-resource --resource-group sample-rg --test-id sample-existing-test-id --keyvault-reference-id null
+    - name: Update autostop criteria.
+      text: |
+        az load test update --load-test-resource sample-alt-resource --resource-group sample-rg --test-id sample-existing-test-id --autostop-error-rate 90 --autostop-time-window 180
+    - name: Update multi-region load configuration.
+      text: |
+        az load test update --load-test-resource sample-alt-resource --resource-group sample-rg --test-id sample-existing-test-id --engine-instances 5 --regionwise-engines eastus=2 westus2=1 eastasia=2
 """
 
 helps[

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="Specify the engine count for each region in the format: region1=engineCount1 region2=engineCount2 .... Use region names in the format accepted by Azure Resource Manager (ARM). Ensure the regions are supported by Azure Load Testing. Multi-region load tests can only target public endpoints.",
+)

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

@@ -0,0 +1,27 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+from dataclasses import dataclass
+
+
+@dataclass
+class LoadTestConfigKeys:
+    DISPLAY_NAME = "displayName"
+    DESCRIPTION = "description"
+    KEYVAULT_REFERENCE_IDENTITY = "keyVaultReferenceIdentity"
+    SUBNET_ID = "subnetId"
+    CERTIFICATES = "certificates"
+    SECRETS = "secrets"
+    ENGINE_INSTANCES = "engineInstances"
+    ENV = "env"
+    PUBLIC_IP_DISABLED = "publicIPDisabled"
+    AUTOSTOP = "autoStop"
+    AUTOSTOP_ERROR_RATE = "errorPercentage"
+    AUTOSTOP_ERROR_RATE_TIME_WINDOW = "timeWindow"
+    FAILURE_CRITERIA = "failureCriteria"
+    REGIONAL_LOADTEST_CONFIG = "regionalLoadTestConfig"
+    REGION = "region"
+    QUICK_START = "quickStartTest"
+    SPLIT_CSV = "splitAllCSVs"

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

@@ -8,6 +8,7 @@
 
 import requests
 import yaml
+from azext_load.data_plane.utils.constants import LoadTestConfigKeys
 from azext_load.data_plane.utils import validators, utils_yaml_config
 from azext_load.vendored_sdks.loadtesting_mgmt import LoadTestMgmtClient
 from azure.cli.core.azclierror import (
@@ -288,43 +289,34 @@ def load_yaml(file_path):
         ) from e
 
 
-def convert_yaml_to_test(data):
+def convert_yaml_to_test(cmd, data):
     new_body = {}
-    if "displayName" in data:
-        new_body["displayName"] = data["displayName"]
-    if "description" in data:
-        new_body["description"] = data["description"]
+    if LoadTestConfigKeys.DISPLAY_NAME in data:
+        new_body["displayName"] = data[LoadTestConfigKeys.DISPLAY_NAME]
+    if LoadTestConfigKeys.DESCRIPTION in data:
+        new_body["description"] = data[LoadTestConfigKeys.DESCRIPTION]
     new_body["keyvaultReferenceIdentityType"] = IdentityType.SystemAssigned
-    if "keyVaultReferenceIdentity" in data:
-        new_body["keyvaultReferenceIdentityId"] = data["keyVaultReferenceIdentity"]
+    if LoadTestConfigKeys.KEYVAULT_REFERENCE_IDENTITY in data:
+        new_body["keyvaultReferenceIdentityId"] = data[LoadTestConfigKeys.KEYVAULT_REFERENCE_IDENTITY]
         new_body["keyvaultReferenceIdentityType"] = IdentityType.UserAssigned
 
-    if "subnetId" in data:
-        new_body["subnetId"] = data["subnetId"]
+    if LoadTestConfigKeys.SUBNET_ID in data:
+        new_body["subnetId"] = data[LoadTestConfigKeys.SUBNET_ID]
 
-    new_body["loadTestConfiguration"] = {}
-    new_body["loadTestConfiguration"]["engineInstances"] = data.get(
-        "engineInstances", 1
-    )
-    if data.get("certificates"):
-        new_body["certificate"] = parse_cert(data.get("certificates"))
-    if data.get("secrets"):
-        new_body["secrets"] = parse_secrets(data.get("secrets"))
-    if data.get("env"):
-        new_body["environmentVariables"] = parse_env(data.get("env"))
-    if data.get("publicIPDisabled"):
-        new_body["publicIPDisabled"] = data.get("publicIPDisabled")
-    # quick test and split csv not supported currently in CLI
-    new_body["loadTestConfiguration"]["quickStartTest"] = False
-    if data.get("quickStartTest"):
-        logger.warning(
-            "Quick start test is not supported currently in CLI. Please use portal to run quick start test"
-        )
-    if data.get("splitAllCSVs") is not None:
-        new_body["loadTestConfiguration"]["splitAllCSVs"] = utils_yaml_config.yaml_parse_splitcsv(data=data)
-    if data.get("failureCriteria"):
+    new_body["loadTestConfiguration"] = utils_yaml_config.yaml_parse_loadtest_configuration(cmd=cmd, data=data)
+
+    if data.get(LoadTestConfigKeys.CERTIFICATES):
+        new_body["certificate"] = parse_cert(data.get(LoadTestConfigKeys.CERTIFICATES))
+    if data.get(LoadTestConfigKeys.SECRETS):
+        new_body["secrets"] = parse_secrets(data.get(LoadTestConfigKeys.SECRETS))
+    if data.get(LoadTestConfigKeys.ENV):
+        new_body["environmentVariables"] = parse_env(data.get(LoadTestConfigKeys.ENV))
+    if data.get(LoadTestConfigKeys.PUBLIC_IP_DISABLED) is not None:
+        new_body["publicIPDisabled"] = data.get(LoadTestConfigKeys.PUBLIC_IP_DISABLED)
+
+    if data.get(LoadTestConfigKeys.FAILURE_CRITERIA):
         new_body["passFailCriteria"] = utils_yaml_config.yaml_parse_failure_criteria(data=data)
-    if data.get("autoStop") is not None:
+    if data.get(LoadTestConfigKeys.AUTOSTOP) is not None:
         new_body["autoStopCriteria"] = utils_yaml_config.yaml_parse_autostop_criteria(data=data)
     logger.debug("Converted yaml to test body: %s", new_body)
     return new_body
@@ -347,6 +339,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."
@@ -428,7 +421,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
 
@@ -502,6 +513,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)."
@@ -558,6 +570,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:
@@ -750,3 +771,16 @@ def upload_files_helper(
         client=client,
         test_id=test_id, yaml_data=yaml_data, test_plan=test_plan,
         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})"
+        )