[load] locust tests support (#8409)

* [load] support multi-region load test configuration

* refactor

* update recordings

* azdev style load

* refactor yaml keys to constants file

* refactor test load test

* update recordings

* azdev style load

* update help text

* update help text

* action on comments

* refactor test regional-load autostop

* check for valid region names client-side

* [load] support advanced url tests

* action on comments

* action on comments

* action on comments

* update recordings

* action on comments

* update test and recordings

* refactor test exception scenarios

* azdev style load

* [load] dashboard reports | debug mode | copy artifacts sas url

* action on comments

* update recordings

* set copy artifacts ut live-only

* change copy-artifacts to get-artifacts

* update recordings

* [load] download-files high scale test update console message

* [load] download-files high scale test update console message

* remove unused import

* update texts

* update texts

* [load] convert test to jmx

* update comment

* update prompt text

* [load] baseline trends | sas download

* unit tests and recordings

* trends aggregate param,unit tests, recordings

* update HISTORY.rst

* remove faulty recording

* action copilot review

Co-authored-by: Copilot <[email protected]>

* undo faulty commit by copilot

* test instruction

* action on comments

* [load] locust test type support

* [load] locust test type support

* azdev mask recordings

* update HISTORY.rst

* copyright in locust script

* azdev scan fix for locust file

* fix history order

* rerun locust test

* run azdev scan

---------

Co-authored-by: Copilot <[email protected]>

Author: mbhardwaj-msft

src/load/HISTORY.rst

@@ -2,6 +2,13 @@
 
 Release History
 ===============
+
+
+1.5.0
+++++++
+* Add support for Locust based load tests.
+
+
 1.4.3
 ++++++
 * Updated the vendored_sdks to include 2024-12-01-preview data plane API version.

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

@@ -7,6 +7,7 @@
 
 import os
 
+from azext_load.data_plane.utils.constants import LoadTestConfigKeys
 from azext_load.data_plane.utils.utils import (
     convert_yaml_to_test,
     create_autostop_criteria_from_args,
@@ -16,6 +17,7 @@
     generate_trends_row,
     get_admin_data_plane_client,
     get_testrun_data_plane_client,
+    infer_test_type_from_test_plan,
     load_yaml,
     upload_file_to_test,
     upload_files_helper,
@@ -71,6 +73,8 @@ def create_test(
     autostop_criteria = create_autostop_criteria_from_args(
         autostop=autostop, error_rate=autostop_error_rate, time_window=autostop_error_rate_time_window)
     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)
         body = create_or_update_test_without_config(
             test_id,
             body,
@@ -91,6 +95,13 @@ def create_test(
     else:
         yaml = load_yaml(load_test_config_file)
         yaml_test_body = convert_yaml_to_test(cmd, yaml)
+        test_type = (
+            test_type or
+            yaml.get(LoadTestConfigKeys.TEST_TYPE) or
+            infer_test_type_from_test_plan(test_plan) or
+            infer_test_type_from_test_plan(yaml.get(LoadTestConfigKeys.TEST_PLAN))
+        )
+        logger.debug("Inferred test type: %s", test_type)
         body = create_or_update_test_with_config(
             test_id,
             body,

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

@@ -43,6 +43,9 @@
     - 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
+    - name: Create a Locust based load test
+      text: |
+        az load test create --test-id sample-test-id --load-test-resource sample-alt-resource --resource-group sample-rg --test-plan ~/resources/sample-locust-file.py --test-type Locust --env LOCUST_HOST="https://azure.microsoft.com" LOCUST_SPAWN_RATE=0.3 LOCUST_RUN_TIME=120 LOCUST_USERS=4
 """
 
 helps[

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

@@ -92,7 +92,7 @@
     validator=validators.validate_test_plan_path,
     options_list=["--test-plan"],
     type=str,
-    help="Reference to the test plan file. If `testType: JMX`: path to the JMeter script. If `testType: URL`: path to the requests JSON file.",
+    help="Reference to the test plan file. If `testType: JMX`: path to the JMeter script. If `testType: URL`: path to the requests JSON file. If `testType: Locust`: path to the Locust test script.",
 )
 
 test_type = CLIArgumentType(

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

@@ -36,11 +36,13 @@ class AllowedMetricNamespaces(str, Enum):
 class AllowedTestTypes(str, Enum):
     JMX = "JMX"
     URL = "URL"
+    LOCUST = "Locust"
 
 
 class AllowedTestPlanFileExtensions(str, Enum):
     JMX = ".jmx"
     URL = ".json"
+    LOCUST = ".py"
 
 
 class AllowedTrendsResponseTimeAggregations(str, Enum):

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, AllowedTestTypes
+from .models import IdentityType, AllowedFileTypes, AllowedTestTypes, AllowedTestPlanFileExtensions
 
 logger = get_logger(__name__)
 
@@ -754,6 +754,19 @@ def upload_zipped_artifacts_helper(
                 )
 
 
+def infer_test_type_from_test_plan(test_plan):
+    if test_plan is None:
+        return None
+    _, file_extension = os.path.splitext(test_plan)
+    if file_extension.casefold() == AllowedTestPlanFileExtensions.JMX.value:
+        return AllowedTestTypes.JMX.value
+    if file_extension.casefold() == AllowedTestPlanFileExtensions.URL.value:
+        return AllowedTestTypes.URL.value
+    if file_extension.casefold() == AllowedTestPlanFileExtensions.LOCUST.value:
+        return AllowedTestTypes.LOCUST.value
+    return None
+
+
 def _evaluate_file_type_for_test_script(test_type, test_plan):
     if test_type == AllowedTestTypes.URL.value:
         _, file_extension = os.path.splitext(test_plan)

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

@@ -121,6 +121,11 @@ class LoadConstants:
 
     HIGH_SCALE_LOAD_TEST_CONFIG_FILE = os.path.join(TEST_RESOURCES_DIR, r"config-high-scale-load.yaml")
 
+    LOCUST_TEST_CONFIG_FILE = os.path.join(TEST_RESOURCES_DIR, r"config-locust.yaml")
+    LOCUST_ENV_VARIABLES = 'LOCUST_HOST="https://www.google.com" LOCUST_SPAWN_RATE=0.3 LOCUST_RUN_TIME=120 LOCUST_USERS=4'
+    LOCUST_TEST_PLAN = os.path.join(TEST_RESOURCES_DIR, r"sample-locust-file.py")
+    LOCUST_TEST_PLAN_FILENAME = "sample-locust-file.py"
+
 
 class LoadTestConstants(LoadConstants):
     # Test IDs for load test commands
@@ -142,6 +147,7 @@ class LoadTestConstants(LoadConstants):
     LOAD_TEST_ADVANCED_URL_ID = "loadtest-advanced-url-case"
     LOAD_TEST_CONVERT_TO_JMX_ID = "loadtest-convert-to-jmx-case"
     LOAD_TEST_BASELINE_TRENDS_ID = "loadtest-baseline-trends-case"
+    LOCUST_LOAD_TEST_ID = "loadtest-locust-case"
 
     INVALID_UPDATE_TEST_ID = "invalid-update-test-case"
     INVALID_PF_TEST_ID = "invalid-pf-test-case"