Merge pull request #1 from mbhardwaj-msft/mbhardwaj/azload-20241202

[load] support advanced url tests

src/load/HISTORY.rst

@@ -6,6 +6,8 @@ Release History
 ++++++
 * 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.
+* Add support for advanced URL test with multiple HTTP request using JSON file. Add `--test-type` argument to 'az load test create' and honor `testType` key in YAML config file.
+
 
 1.3.1
 ++++++

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

@@ -18,7 +18,7 @@
     upload_files_helper,
     create_autostop_criteria_from_args,
 )
-from azure.cli.core.azclierror import InvalidArgumentValueError
+from azure.cli.core.azclierror import InvalidArgumentValueError, FileOperationError
 from azure.core.exceptions import ResourceNotFoundError
 from knack.log import get_logger
 
@@ -31,6 +31,7 @@ def create_test(
     test_id,
     display_name=None,
     test_plan=None,
+    test_type=None,
     resource_group_name=None,
     load_test_config_file=None,
     test_description=None,
@@ -69,6 +70,7 @@ def create_test(
             body,
             display_name=display_name,
             test_description=test_description,
+            test_type=test_type,
             engine_instances=engine_instances,
             env=env,
             secrets=secrets,
@@ -88,6 +90,7 @@ def create_test(
             body,
             yaml_test_body,
             display_name=display_name,
+            test_type=test_type,
             test_description=test_description,
             engine_instances=engine_instances,
             env=env,
@@ -106,8 +109,9 @@ def create_test(
         "Created test with test ID: %s and response obj is %s", test_id, response
     )
     logger.info("Uploading files to test %s", test_id)
+    evaluated_test_type = test_type or yaml_test_body.get("kind") if yaml_test_body else response.get("kind")
     upload_files_helper(
-        client, test_id, yaml, test_plan, load_test_config_file, not custom_no_wait
+        client, test_id, yaml, test_plan, load_test_config_file, not custom_no_wait, evaluated_test_type
     )
     response = client.get_test(test_id)
     logger.info("Upload files to test %s has completed", test_id)
@@ -195,7 +199,7 @@ def update_test(
     )
     logger.info("Uploading files to test %s", test_id)
     upload_files_helper(
-        client, test_id, yaml, test_plan, load_test_config_file, not custom_no_wait
+        client, test_id, yaml, test_plan, load_test_config_file, not custom_no_wait, body.get("kind")
     )
     response = client.get_test(test_id)
     logger.info("Upload files to test %s has completed", test_id)
@@ -409,6 +413,10 @@ def upload_test_file(
     response = upload_file_to_test(
         client, test_id, path, file_type=file_type, wait=not no_wait
     )
+    if not no_wait and response is not None and response.get("validationStatus") == "VALIDATION_FAILURE":
+        raise FileOperationError(
+            f"File upload failed due to validation failure: {response.get('validationFailureDetails')}"
+        )
     logger.debug("Upload test file response: %s", response)
     logger.info("Upload test file completed")
     return response.as_dict()

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

@@ -40,6 +40,9 @@
     - 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
+    - name: Create an advanced URL test with multiple HTTP requests using a JSON file.
+      text: |
+        az load test create --test-id sample-test-id --load-test-resource sample-alt-resource --resource-group sample-rg --test-plan ~/resources/sample-url-requests.json --test-type URL
 """
 
 helps[
@@ -243,4 +246,7 @@
     - name: Upload zipped artifacts to a test.
       text: |
         az load test file upload --test-id sample-test-id --load-test-resource sample-alt-resource --resource-group sample-rg --path ~/Resources/sample-zip.zip --file-type ZIPPED_ARTIFACTS
+    - name: Upload URL requests JSON configuration file to a test which is of type URL.
+      text: |
+        az load test file upload --test-id sample-test-id --load-test-resource sample-alt-resource --resource-group sample-rg --path ~/Resources/sample-url-requests.json --file-type URL_TEST_CONFIG
 """

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

@@ -15,6 +15,7 @@ def load_arguments(self, _):
     with self.argument_context("load test create") as c:
         c.argument("test_id", argtypes.test_id_no_completer)
         c.argument("test_plan", argtypes.test_plan)
+        c.argument("test_type", argtypes.test_type)
         c.argument("display_name", argtypes.test_display_name)
         c.argument("test_description", argtypes.test_description)
         c.argument("env", argtypes.env, help="space-separated environment variables: key[=value] [key[=value] ...].")

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

@@ -92,7 +92,15 @@
     validator=validators.validate_test_plan_path,
     options_list=["--test-plan"],
     type=str,
-    help="Path to the JMeter script.",
+    help="Reference to the test plan file. If `testType: JMX`: path to the JMeter script. If `testType: URL`: path to the requests JSON file.",
+)
+
+test_type = CLIArgumentType(
+    validator=validators.validate_test_type,
+    options_list=["--test-type"],
+    type=str,
+    choices=utils.get_enum_values(models.AllowedTestTypes),
+    help="Type of the load test.",
 )
 
 load_test_config_file = CLIArgumentType(

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

@@ -10,6 +10,8 @@
 class LoadTestConfigKeys:
     DISPLAY_NAME = "displayName"
     DESCRIPTION = "description"
+    TEST_PLAN = "testPlan"
+    TEST_TYPE = "testType"
     KEYVAULT_REFERENCE_IDENTITY = "keyVaultReferenceIdentity"
     SUBNET_ID = "subnetId"
     CERTIFICATES = "certificates"

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

@@ -16,6 +16,8 @@ class AllowedFileTypes(str, Enum):
     JMX_FILE = "JMX_FILE"
     USER_PROPERTIES = "USER_PROPERTIES"
     ZIPPED_ARTIFACTS = "ZIPPED_ARTIFACTS"
+    URL_TEST_CONFIG = "URL_TEST_CONFIG"
+    TEST_SCRIPT = 'TEST_SCRIPT'
 
 
 class AllowedIntervals(str, Enum):
@@ -29,3 +31,13 @@ class AllowedIntervals(str, Enum):
 class AllowedMetricNamespaces(str, Enum):
     LoadTestRunMetrics = "LoadTestRunMetrics"
     EngineHealthMetrics = "EngineHealthMetrics"
+
+
+class AllowedTestTypes(str, Enum):
+    JMX = "JMX"
+    URL = "URL"
+
+
+class AllowedTestPlanFileExtensions(str, Enum):
+    JMX = ".jmx"
+    URL = ".json"

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

@@ -20,7 +20,7 @@
 from azure.mgmt.core.tools import is_valid_resource_id, parse_resource_id
 from knack.log import get_logger
 
-from .models import IdentityType, AllowedFileTypes
+from .models import IdentityType, AllowedFileTypes, AllowedTestTypes
 
 logger = get_logger(__name__)
 
@@ -295,6 +295,8 @@ def convert_yaml_to_test(cmd, data):
         new_body["displayName"] = data[LoadTestConfigKeys.DISPLAY_NAME]
     if LoadTestConfigKeys.DESCRIPTION in data:
         new_body["description"] = data[LoadTestConfigKeys.DESCRIPTION]
+    if LoadTestConfigKeys.TEST_TYPE in data:
+        new_body["kind"] = data[LoadTestConfigKeys.TEST_TYPE]
     new_body["keyvaultReferenceIdentityType"] = IdentityType.SystemAssigned
     if LoadTestConfigKeys.KEYVAULT_REFERENCE_IDENTITY in data:
         new_body["keyvaultReferenceIdentityId"] = data[LoadTestConfigKeys.KEYVAULT_REFERENCE_IDENTITY]
@@ -330,6 +332,7 @@ def create_or_update_test_with_config(
     yaml_test_body,
     display_name=None,
     test_description=None,
+    test_type=None,
     engine_instances=None,
     env=None,
     secrets=None,
@@ -351,6 +354,7 @@ def create_or_update_test_with_config(
         or body.get("displayName")
         or test_id
     )
+    new_body["kind"] = test_type or yaml_test_body.get("kind") or body.get("kind")
 
     test_description = test_description or yaml_test_body.get("description")
     if test_description:
@@ -504,6 +508,7 @@ def create_or_update_test_without_config(
     body,
     display_name=None,
     test_description=None,
+    test_type=None,
     engine_instances=None,
     env=None,
     secrets=None,
@@ -520,6 +525,7 @@ def create_or_update_test_without_config(
     )
     new_body = {}
     new_body["displayName"] = display_name or body.get("displayName") or test_id
+    new_body["kind"] = test_type or body.get("kind")
     test_description = test_description or body.get("description")
     if test_description:
         new_body["description"] = test_description
@@ -723,32 +729,51 @@ def upload_zipped_artifacts_helper(
                 )
 
 
+def _evaluate_file_type_for_test_script(test_type, test_plan):
+    if test_type == AllowedTestTypes.URL.value:
+        _, file_extension = os.path.splitext(test_plan)
+        if file_extension.casefold() == ".json":
+            return AllowedFileTypes.URL_TEST_CONFIG
+        if file_extension.casefold() == ".jmx":
+            return AllowedFileTypes.JMX_FILE
+    return AllowedFileTypes.TEST_SCRIPT
+
+
 def upload_test_plan_helper(
-    client, test_id, yaml_data, test_plan, load_test_config_file, existing_test_files, wait
+    client, test_id, yaml_data, test_plan, load_test_config_file, existing_test_files, wait, test_type
 ):
-    if test_plan is None and yaml_data is not None and yaml_data.get("testPlan"):
-        test_plan = yaml_data.get("testPlan")
+    if test_plan is None and yaml_data is not None and yaml_data.get(LoadTestConfigKeys.TEST_PLAN) is not None:
+        test_plan = yaml_data.get(LoadTestConfigKeys.TEST_PLAN)
     existing_test_plan_files = []
     for file in existing_test_files:
-        if validators.AllowedFileTypes.JMX_FILE.value == file["fileType"]:
+        if (
+            validators.AllowedFileTypes.JMX_FILE.value == file["fileType"] or
+            file["fileType"] == AllowedFileTypes.TEST_SCRIPT.value
+        ):
             existing_test_plan_files.append(file)
     if test_plan:
-        logger.info("Uploading test plan file %s", test_plan)
-        file_response = upload_generic_files_helper(
-            client=client,
-            test_id=test_id, load_test_config_file=load_test_config_file,
-            existing_files=existing_test_plan_files, file_to_upload=test_plan,
-            file_type=validators.AllowedFileTypes.JMX_FILE,
-            wait=wait
-        )
-        if wait and file_response.get("validationStatus") != "VALIDATION_SUCCESS":
-            raise FileOperationError(
-                f"Test plan file {test_plan} is not valid. Please check the file and try again."
+        logger.info("Uploading test plan file %s to test %s of type %s", test_plan, test_id, test_type)
+        file_type = _evaluate_file_type_for_test_script(test_type, test_plan)
+        try:
+            file_response = upload_generic_files_helper(
+                client=client,
+                test_id=test_id, load_test_config_file=load_test_config_file,
+                existing_files=existing_test_files, file_to_upload=test_plan,
+                file_type=file_type,
+                wait=wait
             )
+            if wait and file_response.get("validationStatus") != "VALIDATION_SUCCESS":
+                raise FileOperationError(
+                    f"Test plan file {test_plan} is not valid. Please check the file and try again."
+                )
+        except Exception as e:
+            raise FileOperationError(
+                f"Error occurred while uploading test plan file {test_plan} for test {test_id} of type {test_type}: {e}"
+            ) from e
 
 
 def upload_files_helper(
-    client, test_id, yaml_data, test_plan, load_test_config_file, wait
+    client, test_id, yaml_data, test_plan, load_test_config_file, wait, test_type
 ):
     files = client.list_test_files(test_id)
 
@@ -770,7 +795,8 @@ def upload_files_helper(
     upload_test_plan_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)
+        load_test_config_file=load_test_config_file, existing_test_files=files, wait=wait,
+        test_type=test_type)
 
 
 def validate_engine_data_with_regionwiseload_data(engine_instances, regionwise_engines):

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

@@ -15,7 +15,13 @@
 from knack.log import get_logger
 
 from . import utils
-from .models import AllowedFileTypes, AllowedIntervals, AllowedMetricNamespaces
+from .models import (
+    AllowedFileTypes,
+    AllowedIntervals,
+    AllowedMetricNamespaces,
+    AllowedTestTypes,
+    AllowedTestPlanFileExtensions,
+)
 
 logger = get_logger(__name__)
 
@@ -228,9 +234,25 @@ def validate_test_plan_path(namespace):
     namespace.test_plan = _validate_path(namespace.test_plan, is_dir=False)
 
     _, file_extension = os.path.splitext(namespace.test_plan)
-    if file_extension.casefold() != ".jmx":
+    if file_extension.casefold() not in utils.get_enum_values(AllowedTestPlanFileExtensions):
         raise InvalidArgumentValueError(
-            f"Invalid test plan file extension: {file_extension}. Expected: .jmx"
+            f"Invalid test plan file extension: {file_extension}. "
+            f"Allowed values: {', '.join(AllowedTestPlanFileExtensions)} "
+            f"for {', '.join(utils.get_enum_values(AllowedTestTypes))} test types respectively"
+        )
+
+
+def validate_test_type(namespace):
+    if namespace.test_type is None:
+        return
+    if not isinstance(namespace.test_type, str):
+        raise InvalidArgumentValueError(
+            f"Invalid test-type type: {type(namespace.test_type)}"
+        )
+    allowed_test_types = utils.get_enum_values(AllowedTestTypes)
+    if namespace.test_type not in allowed_test_types:
+        raise InvalidArgumentValueError(
+            f"Invalid test-type value: {namespace.test_type}. Allowed values: {', '.join(allowed_test_types)}"
         )
 
 

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

@@ -58,6 +58,15 @@ class LoadConstants:
     REGIONWISE_ENGINES_INVALID_FORMAT_2 = "germanywestcentral=2 eastus:3"
     REGIONWISE_ENGINES_INVALID_FORMAT_3 = "=2 eastus=3"
 
+    # Constants for Advanced URL Load Tests
+    ADVANCED_URL_FILE_TYPE = "URL_TEST_CONFIG"
+    ADVANCED_URL_TEST_TYPE = "URL"
+    ADVANCED_URL_LOAD_TEST_CONFIG_FILE = os.path.join(TEST_RESOURCES_DIR, r"config-advanced-url.yaml")
+    ADVANCED_TEST_URL_CONFIG_FILE_NAME = "sample-url-requests.json"
+    ADVANCED_TEST_URL_CONFIG_FILE_PATH = os.path.join(TEST_RESOURCES_DIR, r"sample-url-requests.json")
+    ADVANCED_TEST_URL_CONFIG_FILE_UPDATED_NAME = "sample-url-requests-updated.json"
+    ADVANCED_TEST_URL_CONFIG_FILE_UPDATED_PATH = os.path.join(TEST_RESOURCES_DIR, r"sample-url-requests-updated.json")
+
     ENV_VAR_DURATION_NAME = "duration_in_sec"
     ENV_VAR_DURATION_SHORT = "1"
     ENV_VAR_DURATION_LONG = "120"
@@ -126,6 +135,7 @@ class LoadTestConstants(LoadConstants):
     REGIONAL_LOAD_CONFIG_TEST_ID = "regional-load-config-test-case"
     LOAD_TEST_KVREF_ID = "loadtest-kvrefid-case"
     LOAD_TEST_SPLITCSV_ID = "loadtest-splitcsv-case"
+    LOAD_TEST_ADVANCED_URL_ID = "loadtest-advanced-url-case"
 
     INVALID_UPDATE_TEST_ID = "invalid-update-test-case"
     INVALID_PF_TEST_ID = "invalid-pf-test-case"