[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]>
Azure · Feb 27, 2025 · 4e1eb04 · 4e1eb04
1 parent 4ba829a
commit 4e1eb04
Show file tree

Hide file tree

Showing 13 changed files with 3,093 additions and 4 deletions.
diff --git a/src/load/HISTORY.rst b/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.

diff --git a/src/load/azext_load/data_plane/load_test/custom.py b/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,

diff --git a/src/load/azext_load/data_plane/load_test/help.py b/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[

diff --git a/src/load/azext_load/data_plane/utils/argtypes.py b/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(

diff --git a/src/load/azext_load/data_plane/utils/models.py b/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):

diff --git a/src/load/azext_load/data_plane/utils/utils.py b/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)

diff --git a/src/load/azext_load/tests/latest/constants.py b/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"