Add ci test script to test the bundle feature flag in the connectedk8s extension and resolve the ruff check failure

Author: SherryPuffs

src/connectedk8s/azext_connectedk8s/_constants.py

@@ -519,12 +519,18 @@
 allowed_rp_registration_states = ["Registering", "Registered"]
 
 Connected_Cluster_Type = "connectedclusters"
+Arc_Agent_Update_Validator_Job_Label = "agent-update-validator"
 
 Arc_Agentry_Bundle_Feature = "extensionSets"
 Arc_Agentry_Bundle_Feature_Setting = "versionManagedExtensions"
 
 Bundle_Feature_Value_List = ["enabled", "disabled", "preview"]
-Bundle_Extension_Type_List = ["microsoft.iotoperations", "microsoft.extensiondiagnostics", "microsoft.arc.containerstorage", "microsoft.azure.secretstore"]
+Bundle_Extension_Type_List = [
+    "microsoft.iotoperations",
+    "microsoft.extensiondiagnostics",
+    "microsoft.arc.containerstorage",
+    "microsoft.azure.secretstore",
+]
 
 CONST_K8S_EXTENSION_NAME = "k8s-extension"
 CONST_K8S_EXTENSION_CLIENT_FACTORY_MOD_NAME = "azext_k8s_extension._client_factory"

src/connectedk8s/azext_connectedk8s/_help.py

@@ -38,6 +38,10 @@
       text: az connectedk8s connect -g resourceGroupName -n connectedClusterName --enable-oidc-issuer --self-hosted-issuer aksissuerurl
     - name: Onboard a connected kubernetes cluster with azure gateway feature enabled.
       text: az connectedk8s connect -g resourceGroupName -n connectedClusterName --gateway-resource-id gatewayResourceId
+    - name: Onboard a connected kubernetes cluster with the bundling feature enabled.
+      text: az connectedk8s connect -g resourceGroupName -n connectedClusterName --disable-auto-upgrade --config extensionSets.versionManagedExtensions='enabled'
+    - name: Onboard a connected Kubernetes cluster with the bundling feature in preview mode.
+      text: az connectedk8s connect -g resourceGroupName -n connectedClusterName --disable-auto-upgrade --config extensionSets.versionManagedExtensions='preview'
 
 """
 
@@ -61,6 +65,11 @@
       text: az connectedk8s update -g resourceGroupName -n connectedClusterName --gateway-resource-id gatewayResourceId
     - name: Disable the azure gateway feature on a connected kubernetes cluster.
       text: az connectedk8s update -g resourceGroupName -n connectedClusterName --disable-gateway
+    - name: Enable the bundling feature on a connected kubernetes cluster.
+      text: az connectedk8s update -g resourceGroupName -n connectedClusterName --auto-upgrade false --config extensionSets.versionManagedExtensions='enabled'
+    - name: Disable the bundling feature on a connected kubernetes cluster.
+      text: az connectedk8s update -g resourceGroupName -n connectedClusterName --auto-upgrade false --config extensionSets.versionManagedExtensions='disabled'
+
 """
 
 helps["connectedk8s upgrade"] = """

src/connectedk8s/azext_connectedk8s/custom.py

@@ -1542,9 +1542,11 @@ def get_k8s_extension_module(module_name):
     try:
         # adding the installed extension in the path
         from azure.cli.core.extension.operations import add_extension_to_path
+
         add_extension_to_path(consts.CONST_K8S_EXTENSION_NAME)
         # import the extension module
         from importlib import import_module
+
         azext_custom = import_module(module_name)
         return azext_custom
     except ImportError:
@@ -1555,27 +1557,33 @@ def get_k8s_extension_module(module_name):
 
 
 def get_bundle_feature_flag_from_arc_agentry_config(
-    current_arc_agentry_config: list[ArcAgentryConfigurations]
+    current_arc_agentry_config: list[ArcAgentryConfigurations],
 ) -> str | None:
     for agentry_config in current_arc_agentry_config:
-        if agentry_config.feature == consts.Arc_Agentry_Bundle_Feature and \
-           consts.Arc_Agentry_Bundle_Feature_Setting in agentry_config.settings:
-            return agentry_config.settings[consts.Arc_Agentry_Bundle_Feature_Setting].lower()
+        if (
+            agentry_config.feature == consts.Arc_Agentry_Bundle_Feature
+            and consts.Arc_Agentry_Bundle_Feature_Setting in agentry_config.settings
+        ):
+            return agentry_config.settings[
+                consts.Arc_Agentry_Bundle_Feature_Setting
+            ].lower()
     return None
 
 
 def get_bundle_feature_flag_from_configuration_settings(
-    configuration_settings: dict[str, Any]
+    configuration_settings: dict[str, Any],
 ) -> str | None:
     settings = configuration_settings.get(consts.Arc_Agentry_Bundle_Feature, {})
     value = settings.get(consts.Arc_Agentry_Bundle_Feature_Setting)
     return value if value is None else value.lower()
 
 
 def validate_bundle_feature_flag_value(
-    configuration_settings: dict[str, Any]
+    configuration_settings: dict[str, Any],
 ) -> str | None:
-    print(f"Step: {utils.get_utctimestring()}: Validating the bundle feature flag value")
+    print(
+        f"Step: {utils.get_utctimestring()}: Validating the bundle feature flag value"
+    )
     value = get_bundle_feature_flag_from_configuration_settings(configuration_settings)
 
     if value is not None:
@@ -1588,8 +1596,12 @@ def validate_bundle_feature_flag_value(
                 f"'{value}'. Please specify a value from the list: {consts.Bundle_Feature_Value_List}."
             )
 
-        configuration_settings[consts.Arc_Agentry_Bundle_Feature][consts.Arc_Agentry_Bundle_Feature_Setting] = value
-        print(f"Step: {utils.get_utctimestring()}: Setting the bundle feature flag value to '{value}'")
+        configuration_settings[consts.Arc_Agentry_Bundle_Feature][
+            consts.Arc_Agentry_Bundle_Feature_Setting
+        ] = value
+        print(
+            f"Step: {utils.get_utctimestring()}: Setting the bundle feature flag value to '{value}'"
+        )
 
     return value
 
@@ -1598,7 +1610,9 @@ def validate_connect_cluster_bundle_feature_flag_value(
     configuration_settings: dict[str, Any],
     yes: bool = False,
 ):
-    bundle_feature_flag_value = validate_bundle_feature_flag_value(configuration_settings)
+    bundle_feature_flag_value = validate_bundle_feature_flag_value(
+        configuration_settings
+    )
 
     # If the bundle feature flag value is None, skip the validation
     if bundle_feature_flag_value is None:
@@ -1616,7 +1630,7 @@ def validate_connect_cluster_bundle_feature_flag_value(
         logger.warning(
             "Entered %s 'preview' mode. All SLA support is discontinued, and the cluster will remain in 'preview' mode "
             "until it is disconnected from Arc.",
-            consts.Arc_Agentry_Bundle_Feature_Setting
+            consts.Arc_Agentry_Bundle_Feature_Setting,
         )
 
     elif bundle_feature_flag_value == "disabled":
@@ -1635,13 +1649,17 @@ def validate_update_cluster_bundle_feature_flag_value(
     resource_group_name: str,
     cluster_name: str,
 ):
-    bundle_feature_flag_value = validate_bundle_feature_flag_value(configuration_settings)
+    bundle_feature_flag_value = validate_bundle_feature_flag_value(
+        configuration_settings
+    )
 
     # If the bundle feature flag value is None, skip the validation
     if bundle_feature_flag_value is None:
         return
 
-    current_bundle_feature_flag_value = get_bundle_feature_flag_from_arc_agentry_config(current_arc_agentry_config)
+    current_bundle_feature_flag_value = get_bundle_feature_flag_from_arc_agentry_config(
+        current_arc_agentry_config
+    )
 
     if bundle_feature_flag_value == "preview":
         err_msg = (
@@ -1659,8 +1677,11 @@ def validate_update_cluster_bundle_feature_flag_value(
         raise ArgumentUsageError(err_msg)
 
     invalid_transition = (
-        (current_bundle_feature_flag_value == "enabled" and bundle_feature_flag_value == "") or
-        (current_bundle_feature_flag_value == "" and bundle_feature_flag_value == "disabled")
+        current_bundle_feature_flag_value == "enabled"
+        and bundle_feature_flag_value == ""
+    ) or (
+        current_bundle_feature_flag_value == ""
+        and bundle_feature_flag_value == "disabled"
     )
 
     if invalid_transition:
@@ -1672,10 +1693,17 @@ def validate_update_cluster_bundle_feature_flag_value(
         raise ArgumentUsageError(err_msg)
 
     # If the bundle feature flag is set to 'disabled', check if any bundle extensions are installed
-    if current_bundle_feature_flag_value == "enabled" and bundle_feature_flag_value == "disabled":
-        client_factory = get_k8s_extension_module(consts.CONST_K8S_EXTENSION_CLIENT_FACTORY_MOD_NAME)
+    if (
+        current_bundle_feature_flag_value == "enabled"
+        and bundle_feature_flag_value == "disabled"
+    ):
+        client_factory = get_k8s_extension_module(
+            consts.CONST_K8S_EXTENSION_CLIENT_FACTORY_MOD_NAME
+        )
         client = client_factory.cf_k8s_extension_operation(cmd.cli_ctx)
-        k8s_extension_custom_mod = get_k8s_extension_module(consts.CONST_K8S_EXTENSION_CUSTOM_MOD_NAME)
+        k8s_extension_custom_mod = get_k8s_extension_module(
+            consts.CONST_K8S_EXTENSION_CUSTOM_MOD_NAME
+        )
         extensions = k8s_extension_custom_mod.list_k8s_extension(
             client,
             resource_group_name,
@@ -2296,7 +2324,11 @@ def update_connected_cluster(
 
     # Validate and update bundle feature flag value if provided
     validate_update_cluster_bundle_feature_flag_value(
-        cmd, connected_cluster.arc_agentry_configurations, configuration_settings, resource_group_name, cluster_name
+        cmd,
+        connected_cluster.arc_agentry_configurations,
+        configuration_settings,
+        resource_group_name,
+        cluster_name,
     )
 
     arc_agentry_configurations = generate_arc_agent_configuration(
@@ -2869,11 +2901,13 @@ def upgrade_agents(
         ):
             telemetry.set_user_fault()
 
-        namespace = "azure-arc"
-        label_selector = "job-name=agent-update-validator"
+        namespace = consts.Arc_Namespace
+        label_selector = f"job-name={consts.Arc_Agent_Update_Validator_Job_Label}"
 
         # Get the list of pods matching the label
-        pods = api_instance.list_namespaced_pod(namespace=namespace, label_selector=label_selector)
+        pods = api_instance.list_namespaced_pod(
+            namespace=namespace, label_selector=label_selector
+        )
 
         # Extract the terminated message from the container
         if pods.items: