test: az iot ops upgrade integration testing (#470)

azext_edge/tests/edge/init/int/test_init_int.py

@@ -6,7 +6,7 @@
 
 import json
 from os.path import isfile
-from typing import Dict, List, Optional, Union
+from typing import List, Optional
 
 import pytest
 from knack.log import get_logger
@@ -15,7 +15,7 @@
 from azext_edge.edge.util.common import assemble_nargs_to_dict
 
 from ....generators import generate_random_string
-from ....helpers import run
+from ....helpers import process_additional_args, run, strip_quotes
 
 logger = get_logger(__name__)
 
@@ -62,8 +62,8 @@ def init_test_setup(settings, tracked_resources):
         "resourceGroup": settings.env.azext_edge_rg,
         "schemaRegistryId": registry["id"],
         "instanceName": instance_name,
-        "additionalCreateArgs": _strip_quotes(settings.env.azext_edge_create_args),
-        "additionalInitArgs": _strip_quotes(settings.env.azext_edge_init_args),
+        "additionalCreateArgs": strip_quotes(settings.env.azext_edge_create_args),
+        "additionalInitArgs": strip_quotes(settings.env.azext_edge_init_args),
         "continueOnError": settings.env.azext_edge_init_continue_on_error or False,
         "redeployment": settings.env.azext_edge_init_redeployment or False,
     }
@@ -81,9 +81,9 @@ def init_test_setup(settings, tracked_resources):
 @pytest.mark.init_scenario_test
 def test_init_scenario(init_test_setup, tracked_files):
     additional_init_args = init_test_setup["additionalInitArgs"] or ""
-    init_arg_dict = _process_additional_args(additional_init_args)
+    init_arg_dict = process_additional_args(additional_init_args)
     additional_create_args = init_test_setup["additionalCreateArgs"] or ""
-    create_arg_dict = _process_additional_args(additional_create_args)
+    create_arg_dict = process_additional_args(additional_create_args)
     _process_broker_config_file_arg(create_arg_dict, tracked_files)
 
     cluster_name = init_test_setup["clusterName"]
@@ -345,21 +345,6 @@ def assert_trust_config_args(instance_name: str, resource_group: str, trust_sett
         assert issuer["kind"] == trust_args["issuerKind"]
 
 
-def _process_additional_args(additional_args: str) -> Dict[str, Union[str, bool]]:
-    arg_dict = {}
-    for arg in additional_args.split("--")[1:]:
-        arg = arg.strip().split(" ", maxsplit=1)
-        # --simulate-plc vs --desc "potato cluster"
-        arg[0] = arg[0].replace("-", "_")
-        if len(arg) == 1 or arg[1].lower() == "true":
-            arg_dict[arg[0]] = True
-        elif arg[1].lower() == "false":
-            arg_dict[arg[0]] = False
-        else:
-            arg_dict[arg[0]] = arg[1]
-    return arg_dict
-
-
 def _process_broker_config_file_arg(create_arg_dict: dict, tracked_files: List[str]):
     if "broker_config_file" in create_arg_dict:
         broker_config_path = create_arg_dict["broker_config_file"]
@@ -369,14 +354,6 @@ def _process_broker_config_file_arg(create_arg_dict: dict, tracked_files: List[s
                 json.dump(DEFAULT_BROKER_CONFIG, bcf)
 
 
-def _strip_quotes(argument: Optional[str]) -> Optional[str]:
-    if not argument:
-        return argument
-    if argument[0] == argument[-1] and argument[0] in ("'", '"'):
-        argument = argument[1:-1]
-    return argument
-
-
 DEFAULT_BROKER_CONFIG = {
     "advanced": {"encryptInternalTraffic": "Enabled"},
     "cardinality": {

azext_edge/tests/edge/orchestration/test_upgrade_int.py

@@ -0,0 +1,136 @@
+# coding=utf-8
+# ----------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License file in the project root for license information.
+# ----------------------------------------------------------------------------------------------
+
+from copy import deepcopy
+import random
+from typing import Any, Dict, List
+import pytest
+from azext_edge.edge.util import parse_kvp_nargs
+from azext_edge.edge.providers.orchestration.common import EXTENSION_ALIAS_TO_TYPE_MAP
+from ...generators import generate_random_string
+from ...helpers import process_additional_args, run, strip_quotes
+
+EXTENSIONS = list(EXTENSION_ALIAS_TO_TYPE_MAP.keys())
+EXTENSION_TYPE_TO_ALIAS_MAP = {val: key for key, val in EXTENSION_ALIAS_TO_TYPE_MAP.items()}
+
+
+@pytest.fixture
+def upgrade_int_setup(settings):
+    from ...settings import EnvironmentVariables
+
+    settings.add_to_config(EnvironmentVariables.rg.value)
+    settings.add_to_config(EnvironmentVariables.instance.value)
+    settings.add_to_config(EnvironmentVariables.upgrade_args.value)
+
+    if not all([settings.env.azext_edge_instance, settings.env.azext_edge_rg]):
+        raise AssertionError(
+            f"Cannot run init tests without an instance and resource group. Current settings:\n {settings}"
+        )
+
+    yield {
+        "resourceGroup": settings.env.azext_edge_rg,
+        "instanceName": settings.env.azext_edge_instance,
+        "additionalUpgradeArgs": strip_quotes(settings.env.azext_edge_upgrade_args),
+    }
+
+
+@pytest.mark.upgrade_scenario_test
+def test_upgrade(upgrade_int_setup):
+    additional_args = upgrade_int_setup["additionalUpgradeArgs"] or ""
+    resource_group = upgrade_int_setup["resourceGroup"]
+    instance_name = upgrade_int_setup["instanceName"]
+
+    # make tree get us the cluster
+    instance_tree = run(f"az iot ops show -n {instance_name} -g {resource_group} --tree")
+    cluster_name = instance_tree.split("\n", 1)[0].strip()
+
+    cluster_id = run(
+        f"az resource show -n {cluster_name} -g {resource_group} "
+        "--resource-type Microsoft.Kubernetes/connectedClusters"
+    )["id"]
+
+    # get the original extensions and convert it to a map with relevant extensions
+    original_ext_list = get_extensions(cluster_id=cluster_id)
+    original_ext_map = {}
+    for ext in original_ext_list:
+        ext_type = ext["properties"]["extensionType"].lower()
+        if ext_type in EXTENSION_TYPE_TO_ALIAS_MAP:
+            original_ext_map[EXTENSION_TYPE_TO_ALIAS_MAP[ext_type]] = ext
+
+    command = f"az iot ops upgrade -g {resource_group} -n {instance_name} --no-progress -y "
+
+    # run first command with only additional args from input
+    run(f"{command} {additional_args}")
+    assert_extensions(
+        cluster_id=cluster_id,
+        original_ext_map=original_ext_map,
+        additional_args=additional_args
+    )
+    # if additional args present, only run once
+    if additional_args:
+        return
+
+    # run with 2 random config updates
+    upgrade_extensions = random.sample(EXTENSIONS, k=2)
+    for ext in upgrade_extensions:
+        num_config_args = random.choice(range(1, 3))
+        ext_patch = [f"{generate_random_string()}={generate_random_string()}" for _ in range(num_config_args)]
+        ext_arg = f"--{ext}-config {' '.join(ext_patch)} "
+        additional_args += ext_arg
+
+    run(f"{command} {additional_args}")
+    assert_extensions(
+        cluster_id=cluster_id,
+        original_ext_map=original_ext_map,
+        additional_args=additional_args
+    )
+
+
+def assert_extensions(
+    cluster_id: str,
+    original_ext_map: Dict[str, Any],
+    additional_args: str = ""
+):
+    original_ext_map = deepcopy(original_ext_map)
+    additional_args_dict = process_additional_args(additional_args)
+
+    # update the original extensions to the correct value
+    for arg, value in additional_args_dict.items():
+        arg = arg.strip("-").lower()
+        ext, _, operation = arg.partition("_")
+        original_ext = original_ext_map[ext]
+        if operation == "config":
+            parsed_config = parse_kvp_nargs(value.split())
+            original_ext["properties"]["configurationSettings"].update(parsed_config)
+        elif operation == "version":
+            original_ext["properties"]["version"] = value
+        elif operation == "train":
+            original_ext["properties"]["releaseTrain"] = value
+
+    # post upgrade extensions
+    extensions = get_extensions(cluster_id)
+    for extension in extensions:
+        ext_type = extension["properties"]["extensionType"].lower()
+        if ext_type in EXTENSION_TYPE_TO_ALIAS_MAP:
+            ext_type = EXTENSION_TYPE_TO_ALIAS_MAP[ext_type]
+            ext_props = extension["properties"]
+            original_ext_props = original_ext_map[ext_type]["properties"]
+            assert ext_props["configurationSettings"] == original_ext_props["configurationSettings"]
+            assert ext_props["version"] == original_ext_props["version"]
+            assert ext_props["releaseTrain"] == original_ext_props["releaseTrain"]
+
+
+def get_extensions(cluster_id: str) -> List[Dict[str, Any]]:
+    extension_result = run(
+        f"az rest --method GET --url {cluster_id}/providers/"
+        "Microsoft.KubernetesConfiguration/extensions?api-version=2023-05-01"
+    )
+    extensions = extension_result["value"]
+    while extension_result.get("nextLink"):
+        extension_result = run(f"az rest --method GET --url {extension_result['nextLink']}")
+        extensions.extend(extension_result["value"])
+
+    return extensions

azext_edge/tests/helpers.py

@@ -204,6 +204,42 @@ def sort_kubectl_items_by_namespace(
     return sorted_items
 
 
+def process_additional_args(additional_args: str) -> Dict[str, Union[str, bool]]:
+    """
+    Process additional args for init, create, upgrade into dictionaries that can be passed in as kwargs.
+
+    This will transform the args into variable friendly keys (- into _).
+    Flag arguments will be converted to have the boolean value.
+
+    Examples:
+    --simulate-plc -> {"simulate_plc": True}
+    --desc "potato cluster" -> {"desc": "potato cluster"}
+    """
+    arg_dict = {}
+    if not additional_args:
+        return arg_dict
+    for arg in additional_args.split("--")[1:]:
+        arg = arg.strip().split(" ", maxsplit=1)
+        # --simulate-plc vs --desc "potato cluster"
+        arg[0] = arg[0].replace("-", "_")
+        if len(arg) == 1 or arg[1].lower() == "true":
+            arg_dict[arg[0]] = True
+        elif arg[1].lower() == "false":
+            arg_dict[arg[0]] = False
+        else:
+            arg_dict[arg[0]] = arg[1]
+    return arg_dict
+
+
+def strip_quotes(argument: Optional[str]) -> Optional[str]:
+    """Get rid of extra quotes when dealing with pipeline inputs."""
+    if not argument:
+        return argument
+    if argument[0] == argument[-1] and argument[0] in ("'", '"'):
+        argument = argument[1:-1]
+    return argument
+
+
 def generate_ops_resource(segments: int = 1) -> IoTOperationsResource:
     resource_id = ""
     for _ in range(segments):

azext_edge/tests/settings.py

@@ -25,6 +25,7 @@ class EnvironmentVariables(Enum):
     aio_cleanup = "azext_edge_aio_cleanup"
     init_continue_on_error = "azext_edge_init_continue_on_error"
     init_redeployment = "azext_edge_init_redeployment"
+    upgrade_args = "azext_edge_upgrade_args"
 
 
 class Setting(object):

pytest.ini.example

@@ -20,3 +20,4 @@ env =
 
 markers =
     init_scenario_test: mark tests that will run az iot ops init
+    upgrade_scenario_test: mark tests that will run az iot ops upgrade