[Azure Load Testing] Adding disable-public-ip argument in load extension

src/load/HISTORY.rst

@@ -2,6 +2,10 @@
 
 Release History
 ===============
+1.2.0
+++++++
+* Added support for disable public IP in test creation and update. This can be done by using --disable-public-ip argument in 'az load test create' and 'az load test update' commands.
+
 1.1.1
 ++++++
 * Fix empty response object on CLI when using 'az load test file upload'

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

@@ -40,6 +40,7 @@ def create_test(
     key_vault_reference_identity=None,
     subnet_id=None,
     split_csv=None,
+    disable_public_ip=None,
     custom_no_wait=False,
 ):
     client = get_admin_data_plane_client(cmd, load_test_resource, resource_group_name)
@@ -70,6 +71,7 @@ def create_test(
             key_vault_reference_identity=key_vault_reference_identity,
             subnet_id=subnet_id,
             split_csv=split_csv,
+            disable_public_ip=disable_public_ip,
         )
     else:
         yaml = load_yaml(load_test_config_file)
@@ -87,6 +89,7 @@ def create_test(
             key_vault_reference_identity=key_vault_reference_identity,
             subnet_id=subnet_id,
             split_csv=split_csv,
+            disable_public_ip=disable_public_ip,
         )
     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)
@@ -119,6 +122,7 @@ def update_test(
     key_vault_reference_identity=None,
     subnet_id=None,
     split_csv=None,
+    disable_public_ip=None,
     custom_no_wait=False,
 ):
     client = get_admin_data_plane_client(cmd, load_test_resource, resource_group_name)
@@ -148,6 +152,7 @@ def update_test(
             key_vault_reference_identity=key_vault_reference_identity,
             subnet_id=subnet_id,
             split_csv=split_csv,
+            disable_public_ip=disable_public_ip,
         )
     else:
         body = create_or_update_test_without_config(
@@ -162,6 +167,7 @@ def update_test(
             key_vault_reference_identity=key_vault_reference_identity,
             subnet_id=subnet_id,
             split_csv=split_csv,
+            disable_public_ip=disable_public_ip,
         )
     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/params.py

@@ -28,6 +28,7 @@ def load_arguments(self, _):
         )
         c.argument("engine_instances", argtypes.engine_instances)
         c.argument("custom_no_wait", argtypes.custom_no_wait)
+        c.argument("disable_public_ip", argtypes.disable_public_ip)
 
     with self.argument_context("load test update") as c:
         c.argument("load_test_config_file", argtypes.load_test_config_file)
@@ -44,6 +45,7 @@ def load_arguments(self, _):
         c.argument("subnet_id", argtypes.subnet_id)
         c.argument("split_csv", argtypes.split_csv)
         c.argument("custom_no_wait", argtypes.custom_no_wait)
+        c.argument("disable_public_ip", argtypes.disable_public_ip)
 
     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

@@ -34,6 +34,14 @@
     help="Do not wait for the long-running operation to finish.",
 )
 
+disable_public_ip = CLIArgumentType(
+    validator=validators.validate_disable_public_ip,
+    options_list=["--disable-public-ip"],
+    type=str,
+    help="Disable the deployment of a public IP address, load balancer, and network security group while testing a private endpoint.",
+)
+
+
 force = CLIArgumentType(
     options_list=["--force"],
     action="store_true",

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

@@ -301,7 +301,8 @@ def convert_yaml_to_test(data):
         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"):
@@ -365,6 +366,7 @@ def create_or_update_test_with_config(
     key_vault_reference_identity=None,
     subnet_id=None,
     split_csv=None,
+    disable_public_ip=None,
 ):
     logger.info(
         "Creating a request body for create or update test using config and parameters."
@@ -397,6 +399,10 @@ def create_or_update_test_with_config(
             new_body["keyvaultReferenceIdentityType"] = IdentityType.SystemAssigned
             new_body.pop("keyvaultReferenceIdentityId")
     subnet_id = subnet_id or yaml_test_body.get("subnetId")
+    if disable_public_ip is not None:
+        new_body["publicIPDisabled"] = disable_public_ip
+    else:
+        new_body["publicIPDisabled"] = yaml_test_body.get("publicIPDisabled", False)
     if subnet_id:
         if subnet_id.casefold() in ["null", ""]:
             new_body["subnetId"] = None
@@ -485,6 +491,7 @@ def create_or_update_test_without_config(
     key_vault_reference_identity=None,
     subnet_id=None,
     split_csv=None,
+    disable_public_ip=None,
 ):
     logger.info(
         "Creating a request body for test using parameters and old test body (in case of update)."
@@ -549,6 +556,8 @@ def create_or_update_test_without_config(
         new_body["loadTestConfiguration"]["splitAllCSVs"] = body[
             "loadTestConfiguration"
         ]["splitAllCSVs"]
+    if disable_public_ip is not None:
+        new_body["publicIPDisabled"] = disable_public_ip
     logger.debug("Request body for create or update test: %s", new_body)
     return new_body
 

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

@@ -382,3 +382,23 @@ def validate_split_csv(namespace):
         namespace.split_csv = True
     else:
         namespace.split_csv = False
+
+
+def validate_disable_public_ip(namespace):
+    if namespace.disable_public_ip is None:
+        return
+    if not isinstance(namespace.disable_public_ip, str):
+        raise InvalidArgumentValueError(
+            f"Invalid disable-public-ip type: {type(namespace.disable_public_ip)}"
+        )
+    if namespace.disable_public_ip.casefold() not in [
+        "true",
+        "false",
+    ]:
+        raise InvalidArgumentValueError(
+            f"Invalid disable-public-ip value: {namespace.disable_public_ip}. Allowed values: true, false"
+        )
+    if namespace.disable_public_ip.casefold() in ["true"]:
+        namespace.disable_public_ip = True
+    else:
+        namespace.disable_public_ip = False

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

@@ -11,6 +11,8 @@
 class LoadConstants:
     # Test Plan constants
     LOAD_TEST_CONFIG_FILE = os.path.join(TEST_RESOURCES_DIR, r"config.yaml")
+    LOAD_TEST_CONFIG_FILE_PUBLIC_IP_DISABLED_FALSE = os.path.join(TEST_RESOURCES_DIR, r"config-disable-public-ip-false.yaml")
+    LOAD_TEST_CONFIG_FILE_PUBLIC_IP_DISABLED_TRUE = os.path.join(TEST_RESOURCES_DIR, r"config-disable-public-ip-true.yaml")
     INVALID_LOAD_TEST_CONFIG_FILE = os.path.join(
         TEST_RESOURCES_DIR, r"invalid-config.yaml"
     )
@@ -48,6 +50,9 @@ class LoadConstants:
     SPLIT_CSV_TRUE = "true"
     SPLIT_CSV_FALSE = "false"
 
+    DISABLE_PUBLIC_IP_TRUE = "true"
+    DISABLE_PUBLIC_IP_FALSE = "false"
+
     INVALID_SUBNET_ID = r"/subscriptions/invalid/resource/id"
     KEYVAULT_REFERENCE_ID = r"/subscriptions/00000000-0000-0000-0000-000000000000/resourcegroups/sample-rg/providers/microsoft.managedidentity/userassignedidentities/sample-mi"
     # App Component constants
@@ -71,6 +76,7 @@ class LoadConstants:
 class LoadTestConstants(LoadConstants):
     # Test IDs for load test commands
     UPDATE_WITH_CONFIG_TEST_ID = "update-with-config-test-case"
+    CREATE_AND_UPDATE_VNET_TEST_ID = "create-update-vnet-test-case"
     DELETE_TEST_ID = "delete-test-case"
     CREATE_TEST_ID = "create-test-case"
     UPDATE_TEST_ID = "update-test-case"
@@ -85,6 +91,7 @@ class LoadTestConstants(LoadConstants):
     INVALID_UPDATE_TEST_ID = "invalid-update-test-case"
     INVALID_PF_TEST_ID = "invalid-pf-test-case"
     INVALID_ZIP_COUNT_TEST_ID = "invalid-zip-count-test-case"
+    INVALID_DISABLED_PUBLIC_IP_TEST_ID = "invalid-disable-public-ip-test-case"
 
     DESCRIPTION = r"Sample_test_description"
     DISPLAY_NAME = r"Sample_test_display_name"