{AKS} Stabilize aks-preview live tests

src/aks-preview/azext_aks_preview/_helpers.py

@@ -2,6 +2,8 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
+from collections.abc import MutableMapping
+
 import errno
 import os
 import platform
@@ -44,6 +46,28 @@
 allowed_extensions = ["microsoft.dataprotection.kubernetes"]
 
 
+def reset_agentpool_to_name_and_mode(agentpool, mode):
+    """Remove all agent pool fields except the resource name and pool mode."""
+    name = agentpool.name
+    properties = getattr(agentpool, "properties", None)
+    if isinstance(properties, MutableMapping):
+        properties.clear()
+        properties["mode"] = mode
+        agentpool.clear()
+        agentpool["name"] = name
+        agentpool["properties"] = properties
+    elif isinstance(agentpool, MutableMapping):
+        agentpool.clear()
+        agentpool["name"] = name
+        agentpool["mode"] = mode
+    else:
+        agentpool.mode = mode
+    for attr in list(vars(agentpool)):
+        if attr not in ("name", "mode") and not attr.startswith("_") and hasattr(agentpool, attr):
+            setattr(agentpool, attr, None)
+    return agentpool
+
+
 def which(binary):
     path_var = os.getenv('PATH')
     if platform.system() == 'Windows':

src/aks-preview/azext_aks_preview/_loadbalancer.py

@@ -172,8 +172,11 @@ def configure_load_balancer_profile(
                 profile.managed_outbound_i_ps = (
                     ManagedClusterLoadBalancerProfileManagedOutboundIPs()
                 )
-            if managed_outbound_ip_count is not None:
-                profile.managed_outbound_i_ps.count = managed_outbound_ip_count
+            profile.managed_outbound_i_ps.count = (
+                managed_outbound_ip_count
+                if managed_outbound_ip_count is not None
+                else profile.managed_outbound_i_ps.count or 1
+            )
             if managed_outbound_ipv6_count is not None:
                 profile.managed_outbound_i_ps.count_ipv6 = managed_outbound_ipv6_count
         else:

src/aks-preview/azext_aks_preview/agentpool_decorator.py

@@ -57,6 +57,7 @@
     get_nodepool_snapshot_by_snapshot_id,
     filter_hard_taints,
     process_dns_overrides,
+    reset_agentpool_to_name_and_mode,
 )
 
 logger = get_logger(__name__)
@@ -1615,14 +1616,9 @@ def set_up_managed_system_mode(self, agentpool: AgentPool) -> AgentPool:
             if agentpool is None:
                 raise CLIInternalError("agentpool cannot be None for ManagedSystem mode")
 
-            # Instead of creating a new instance, modify the existing one
-            # Keep name and set mode to ManagedSystem
-            agentpool.mode = CONST_NODEPOOL_MODE_MANAGEDSYSTEM
-            # Make sure all other attributes are None
-            for attr in vars(agentpool):
-                if attr != 'name' and attr != 'mode' and not attr.startswith('_'):
-                    if hasattr(agentpool, attr):
-                        setattr(agentpool, attr, None)
+            agentpool = reset_agentpool_to_name_and_mode(
+                agentpool, CONST_NODEPOOL_MODE_MANAGEDSYSTEM
+            )
 
         return agentpool
 
@@ -1636,28 +1632,9 @@ def set_up_machines_mode(self, agentpool: AgentPool) -> AgentPool:
 
         mode = self.context.get_mode()
         if mode == CONST_NODEPOOL_MODE_MACHINES:
-            agentpool.mode = CONST_NODEPOOL_MODE_MACHINES
-            # Make sure all other attributes are None
-            # Check properties sub-model first (AgentPool), then flat fields (ManagedClusterAgentPoolProfile)
-            props = getattr(agentpool, 'properties', None)
-            rest_fields = getattr(props, '_attr_to_rest_field', None) if props is not None else None
-            if rest_fields is not None:
-                target, fields = props, rest_fields
-            else:
-                rest_fields = getattr(agentpool, '_attr_to_rest_field', None)
-                if rest_fields is not None and 'mode' in rest_fields:
-                    target, fields = agentpool, rest_fields
-                else:
-                    target, fields = None, None
-            if target is not None:
-                for attr in list(fields.keys()):
-                    if attr not in ('name', 'mode'):
-                        setattr(agentpool, attr, None)
-            else:
-                for attr in vars(agentpool):
-                    if attr != 'name' and attr != 'mode' and not attr.startswith('_'):
-                        if hasattr(agentpool, attr):
-                            setattr(agentpool, attr, None)
+            agentpool = reset_agentpool_to_name_and_mode(
+                agentpool, CONST_NODEPOOL_MODE_MACHINES
+            )
 
         return agentpool
 
@@ -2119,12 +2096,9 @@ def update_agentpool_profile_preview(self, agentpools: List[AgentPool] = None) -
 
         # Check if agentpool is in ManagedSystem mode and handle special case
         if agentpool.mode == CONST_NODEPOOL_MODE_MANAGEDSYSTEM:
-            # Make sure all other attributes are None
-            for attr in vars(agentpool):
-                if attr != 'name' and attr != 'mode' and not attr.startswith('_'):
-                    if hasattr(agentpool, attr):
-                        setattr(agentpool, attr, None)
-            return agentpool
+            return reset_agentpool_to_name_and_mode(
+                agentpool, CONST_NODEPOOL_MODE_MANAGEDSYSTEM
+            )
 
         # update network profile
         agentpool = self.update_network_profile(agentpool)

src/aks-preview/azext_aks_preview/aks_diagnostics.py

@@ -3,6 +3,8 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 
+from collections.abc import Mapping
+
 import datetime
 import json
 import os
@@ -45,6 +47,16 @@ class ClusterFeatures(Flag):
     WIN_HPC = auto()
 
 
+def _get_storage_account_key(storage_account_keys):
+    keys = (
+        storage_account_keys["keys"]
+        if isinstance(storage_account_keys, Mapping)
+        else storage_account_keys.keys
+    )
+    first_key = keys[0]
+    return first_key["value"] if isinstance(first_key, Mapping) else first_key.value
+
+
 # pylint: disable=line-too-long
 def aks_kollect_cmd(cmd,    # pylint: disable=too-many-statements,too-many-locals
                     client,
@@ -99,17 +111,18 @@ def aks_kollect_cmd(cmd,    # pylint: disable=too-many-statements,too-many-local
             cmd.cli_ctx, parsed_storage_account['subscription'])
         storage_account_keys = storage_client.storage_accounts.list_keys(parsed_storage_account['resource_group'],
                                                                          storage_account_name)
+        storage_account_key = _get_storage_account_key(storage_account_keys)
 
         t_generate_blob_service_sas = get_sdk(cmd.cli_ctx, ResourceType.DATA_STORAGE_BLOB, '#generate_account_sas')
 
         sas_token = t_generate_blob_service_sas(storage_account_name,
-                                                storage_account_keys.keys[0].value,
+                                                storage_account_key,
                                                 resource_types='sco',
                                                 permission='rwdlacup',
                                                 expiry=datetime.datetime.utcnow() + datetime.timedelta(days=1))
 
         readonly_sas_token = t_generate_blob_service_sas(storage_account_name,
-                                                         storage_account_keys.keys[0].value,
+                                                         storage_account_key,
                                                          resource_types='sco',
                                                          permission='rl',
                                                          expiry=datetime.datetime.utcnow() + datetime.timedelta(days=1))

src/aks-preview/azext_aks_preview/managed_cluster_decorator.py

@@ -68,6 +68,7 @@
     get_cluster_snapshot_by_snapshot_id,
     get_monitoring_addon_key,
     filter_hard_taints,
+    reset_agentpool_to_name_and_mode,
 )
 from azext_aks_preview._loadbalancer import create_load_balancer_profile
 from azext_aks_preview._loadbalancer import (
@@ -6200,26 +6201,9 @@ def update_managed_system_pools(self, mc: ManagedCluster) -> ManagedCluster:
             # Check if agentpool is in ManagedSystem mode and handle special case
             if agentpool.mode != CONST_NODEPOOL_MODE_MANAGEDSYSTEM:
                 continue
-            # Make sure all other attributes are None
-            # Check properties sub-model first (AgentPool), then flat fields (ManagedClusterAgentPoolProfile)
-            props = getattr(agentpool, 'properties', None)
-            rest_fields = getattr(props, '_attr_to_rest_field', None) if props is not None else None
-            if rest_fields is not None:
-                target, fields = props, rest_fields
-            else:
-                rest_fields = getattr(agentpool, '_attr_to_rest_field', None)
-                if rest_fields is not None and 'mode' in rest_fields:
-                    target, fields = agentpool, rest_fields
-                else:
-                    target, fields = None, None
-            if target is not None:
-                for attr in list(fields.keys()):
-                    if attr not in ('name', 'mode'):
-                        setattr(agentpool, attr, None)
-            else:
-                for attr in vars(agentpool):
-                    if attr not in ('name', 'mode') and not attr.startswith('_') and hasattr(agentpool, attr):
-                        setattr(agentpool, attr, None)
+            reset_agentpool_to_name_and_mode(
+                agentpool, CONST_NODEPOOL_MODE_MANAGEDSYSTEM
+            )
         return mc
 
     def init_models(self) -> None:

src/aks-preview/azext_aks_preview/tests/latest/test_agentpool_decorator.py

@@ -2054,12 +2054,12 @@ def common_set_up_managed_system_mode(self):
         # Verify that name is preserved
         self.assertEqual(dec_agentpool_1.name, original_name)
 
-        # Verify that all other properties are reset to None
-        for attr_name in vars(dec_agentpool_1):
-            if attr_name not in ['name', 'mode'] and not attr_name.startswith('_'):
-                attr_value = getattr(dec_agentpool_1, attr_name)
-                self.assertIsNone(attr_value,
-                    f"Attribute '{attr_name}' should be None but was '{attr_value}'")
+        self.assertIsNone(dec_agentpool_1.count)
+        self.assertIsNone(dec_agentpool_1.vm_size)
+        self.assertIsNone(dec_agentpool_1.os_type)
+        self.assertIsNone(dec_agentpool_1.enable_auto_scaling)
+        self.assertIsNone(dec_agentpool_1.min_count)
+        self.assertIsNone(dec_agentpool_1.max_count)
 
         # Test case 2: mode is not ManagedSystem - should return agentpool unchanged
         dec_2 = AKSPreviewAgentPoolAddDecorator(
@@ -2138,11 +2138,12 @@ def common_set_up_machines_mode(self):
         dec_agentpool_1 = dec_1.set_up_machines_mode(agentpool_1)
         self.assertEqual(dec_agentpool_1.name, original_name)
         self.assertEqual(dec_agentpool_1.mode, CONST_NODEPOOL_MODE_MACHINES)
-        for attr_name in vars(dec_agentpool_1):
-            if attr_name not in ['name', 'mode'] and not attr_name.startswith('_'):
-                attr_value = getattr(dec_agentpool_1, attr_name)
-                self.assertIsNone(attr_value,
-                    f"Attribute '{attr_name}' should be None but was '{attr_value}'")
+        self.assertIsNone(dec_agentpool_1.count)
+        self.assertIsNone(dec_agentpool_1.vm_size)
+        self.assertIsNone(dec_agentpool_1.os_type)
+        self.assertIsNone(dec_agentpool_1.enable_auto_scaling)
+        self.assertIsNone(dec_agentpool_1.min_count)
+        self.assertIsNone(dec_agentpool_1.max_count)
 
     def common_construct_agentpool_profile_preview_with_managed_system_mode(self):
         """Test that construct_agentpool_profile_preview properly handles ManagedSystem mode"""