[load] support for advanced urls test

Author: mbhardwaj-msft

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

@@ -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,
@@ -68,6 +69,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,
@@ -87,6 +89,7 @@ def create_test(
             yaml_test_body,
             display_name=display_name,
             test_description=test_description,
+            test_type=test_type,
             engine_instances=engine_instances,
             env=env,
             secrets=secrets,
@@ -104,7 +107,7 @@ def create_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, test_type or yaml_test_body.get("kind") if yaml_test_body else None
     )
     response = client.get_test(test_id)
     logger.info("Upload files to test %s has completed", test_id)
@@ -189,7 +192,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)

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

@@ -95,6 +95,13 @@
     help="Path to the JMeter script.",
 )
 
+test_type = CLIArgumentType(
+    validator=validators.validate_test_type,
+    options_list=["--test-type"],
+    type=str,
+    help=f"Type of the load test. Allowed values: {', '.join(utils.get_enum_values(models.AllowedTestTypes))}.",
+)
+
 load_test_config_file = CLIArgumentType(
     validator=validators.validate_load_test_config_file,
     options_list=["--load-test-config-file"],

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,7 @@ class AllowedIntervals(str, Enum):
 class AllowedMetricNamespaces(str, Enum):
     LoadTestRunMetrics = "LoadTestRunMetrics"
     EngineHealthMetrics = "EngineHealthMetrics"
+
+class AllowedTestTypes(str, Enum):
+    JMX = "JMX"
+    URL = "URL"

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

@@ -19,7 +19,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__)
 
@@ -180,7 +180,7 @@ def upload_file_to_test(client, test_id, file_path, file_type=None, wait=False):
     # pylint: disable-next=protected-access
     file_path = validators._validate_path(file_path, is_dir=False)
     # pylint: disable-next=protected-access
-    validators._validate_file_stats(file_path, file_type)
+    validators._validate_file_stats(file_path, file_type)    
     with open(file_path, "rb") as file:
         upload_poller = client.begin_upload_test_file(
             test_id,
@@ -294,6 +294,8 @@ def convert_yaml_to_test(data):
         new_body["displayName"] = data["displayName"]
     if "description" in data:
         new_body["description"] = data["description"]
+    if "testType" in data:
+        new_body["kind"] = data["testType"]
     new_body["keyvaultReferenceIdentityType"] = IdentityType.SystemAssigned
     if "keyVaultReferenceIdentity" in data:
         new_body["keyvaultReferenceIdentityId"] = data["keyVaultReferenceIdentity"]
@@ -338,6 +340,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,
@@ -358,6 +361,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:
@@ -493,6 +497,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,
@@ -508,6 +513,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
@@ -702,32 +708,49 @@ 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 == ".json":
+            return AllowedFileTypes.URL_TEST_CONFIG
+        if file_extension == ".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")
     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."
+        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}: {str(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)
 
@@ -749,4 +772,5 @@ 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)

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

@@ -14,7 +14,7 @@
 from knack.log import get_logger
 
 from . import utils
-from .models import AllowedFileTypes, AllowedIntervals, AllowedMetricNamespaces
+from .models import AllowedFileTypes, AllowedIntervals, AllowedMetricNamespaces, AllowedTestTypes
 
 logger = get_logger(__name__)
 
@@ -233,6 +233,20 @@ def validate_test_plan_path(namespace):
         )
 
 
+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)}"
+        )
+
+
 def _validate_path(path, is_dir=False):
     logger.info("path: %s", path)
     if not isinstance(path, str):