add unittest

Author: Hao Yuan

src/aks-preview/azext_aks_preview/agentpool_decorator.py

@@ -17,6 +17,7 @@
     AKSAgentPoolUpdateDecorator,
 )
 from azure.cli.core.azclierror import (
+    CLIInternalError,
     InvalidArgumentValueError,
     MutuallyExclusiveArgumentError,
 )
@@ -1079,6 +1080,10 @@ def set_up_managed_system_mode(self, agentpool: AgentPool) -> AgentPool:
         """
         mode = self.context.raw_param.get("mode")
         if mode == CONST_NODEPOOL_MODE_MANAGEDSYSTEM:
+            # Raise error if agentpool is None
+            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
             name = agentpool.name

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

@@ -34,6 +34,7 @@
     CONST_DEFAULT_WINDOWS_VMS_VM_SIZE,
     CONST_GPU_DRIVER_INSTALL,
     CONST_GPU_DRIVER_NONE,
+    CONST_NODEPOOL_MODE_MANAGEDSYSTEM,
 )
 from azure.cli.command_modules.acs.agentpool_decorator import AKSAgentPoolParamDict
 from azure.cli.command_modules.acs.tests.latest.mocks import (
@@ -424,7 +425,7 @@ def common_get_driver_type(self):
                 driver_type="CUDA"
             )
         )
-        
+
         ctx_1.attach_agentpool(agentpool_1)
         self.assertEqual(ctx_1.get_driver_type(), "CUDA")
 
@@ -645,7 +646,7 @@ def common_get_disable_vtpm(self):
         )
         ctx_1.attach_agentpool(agentpool_1)
         self.assertEqual(ctx_1.get_disable_vtpm(), True)
-    
+
     def common_get_enable_fips_image(self):
         # default
         ctx_1 = AKSPreviewAgentPoolContext(
@@ -659,7 +660,7 @@ def common_get_enable_fips_image(self):
         agentpool = self.create_initialized_agentpool_instance(enable_fips=True)
         ctx_1.attach_agentpool(agentpool)
         self.assertEqual(ctx_1.get_enable_fips_image(), True)
-        
+
         # default
         ctx_2 = AKSPreviewAgentPoolContext(
             self.cmd,
@@ -782,7 +783,7 @@ def common_get_node_vm_size(self):
                 ctx_4.get_node_vm_size()
         else:
             self.assertEqual(ctx_4.get_node_vm_size(), CONST_DEFAULT_WINDOWS_NODE_VM_SIZE)
-        
+
         # if --node-vm-size is not specified, but --sku automatic is explicitly specified
         ctx_5 = AKSPreviewAgentPoolContext(
             self.cmd,
@@ -928,10 +929,10 @@ def test_get_enable_vtpm(self):
 
     def test_get_disable_vtpm(self):
         self.common_get_disable_vtpm()
-    
+
     def common_get_enable_fips_image(self):
         self.common_get_enable_fips_image()
-    
+
     def common_get_disable_fips_image(self):
         self.common_get_disable_fips_image()
 
@@ -958,6 +959,7 @@ def setUp(self):
         self.models = AKSPreviewAgentPoolModels(
             self.cmd, self.resource_type, self.agentpool_decorator_mode
         )
+        self.client = MockClient()
 
     def test_get_zones(self):
         self.common_get_zones()
@@ -1000,7 +1002,7 @@ def test_get_disable_secure_boot(self):
 
     def test_get_enable_vtpm(self):
         self.common_get_enable_vtpm()
-    
+
     def common_get_enable_fips_image(self):
         self.common_get_enable_fips_image()
 
@@ -1013,6 +1015,80 @@ def test_get_gateway_prefix_size(self):
     def test_get_vm_sizes(self):
         self.common_get_vm_sizes()
 
+    def test_construct_agentpool_profile_preview(self):
+        import inspect
+
+        from azext_aks_preview.custom import aks_create
+
+        optional_params = {}
+        positional_params = []
+        for _, v in inspect.signature(aks_create).parameters.items():
+            if v.default != v.empty:
+                optional_params[v.name] = v.default
+            else:
+                positional_params.append(v.name)
+        ground_truth_positional_params = [
+            "cmd",
+            "client",
+            "resource_group_name",
+            "name",
+            "ssh_key_value",
+        ]
+        self.assertEqual(positional_params, ground_truth_positional_params)
+
+        # prepare a dictionary of default parameters
+        raw_param_dict = {
+            "resource_group_name": "test_rg_name",
+            "name": "test_cluster_name",
+            "ssh_key_value": None,
+        }
+        raw_param_dict.update(optional_params)
+
+        # default value in `aks nodepool add`
+        dec_1 = AKSPreviewAgentPoolAddDecorator(
+            self.cmd,
+            self.client,
+            raw_param_dict,
+            self.resource_type,
+            self.agentpool_decorator_mode,
+        )
+
+        with patch(
+            "azure.cli.command_modules.acs.agentpool_decorator.cf_agent_pools",
+            return_value=Mock(list=Mock(return_value=[])),
+        ):
+            dec_agentpool_1 = dec_1.construct_agentpool_profile_preview()
+
+        upgrade_settings_1 = self.models.AgentPoolUpgradeSettings()
+        # CLI will create sshAccess=localuser by default
+        ground_truth_security_profile = self.models.AgentPoolSecurityProfile()
+        ground_truth_security_profile.ssh_access = CONST_SSH_ACCESS_LOCALUSER
+        ground_truth_agentpool_1 = self.create_initialized_agentpool_instance(
+            nodepool_name="nodepool1",
+            orchestrator_version="",
+            vm_size=CONST_DEFAULT_NODE_VM_SIZE,
+            os_type=CONST_DEFAULT_NODE_OS_TYPE,
+            enable_node_public_ip=False,
+            enable_auto_scaling=False,
+            count=3,
+            node_taints=[],
+            node_initialization_taints=[],
+            os_disk_size_gb=0,
+            upgrade_settings=upgrade_settings_1,
+            type=CONST_VIRTUAL_MACHINE_SCALE_SETS,
+            enable_encryption_at_host=False,
+            enable_ultra_ssd=False,
+            enable_fips=False,
+            mode=CONST_NODEPOOL_MODE_SYSTEM,
+            workload_runtime=CONST_WORKLOAD_RUNTIME_OCI_CONTAINER,
+            enable_custom_ca_trust=False,
+            network_profile=self.models.AgentPoolNetworkProfile(),
+            security_profile=ground_truth_security_profile,
+        )
+        self.assertEqual(dec_agentpool_1, ground_truth_agentpool_1)
+
+        dec_1.context.raw_param.print_usage_statistics()
+
 
 class AKSPreviewAgentPoolAddDecoratorCommonTestCase(unittest.TestCase):
     def _remove_defaults_in_agentpool(self, agentpool):
@@ -1333,7 +1409,7 @@ def common_set_up_virtual_machines_profile(self):
             )
         )
         self.assertEqual(dec_agentpool_1, ground_truth_agentpool_1)
-        
+
         dec_2 = AKSPreviewAgentPoolAddDecorator(
             self.cmd,
             self.client,
@@ -1347,6 +1423,142 @@ def common_set_up_virtual_machines_profile(self):
         with self.assertRaises(InvalidArgumentValueError):
             dec_2.set_up_virtual_machines_profile(agentpool_2)
 
+    def common_set_up_managed_system_mode(self):
+        """Test the set_up_managed_system_mode method in AKSPreviewAgentPoolAddDecorator"""
+
+        # Test case 1: mode is ManagedSystem - should reset all properties except name and mode
+        dec_1 = AKSPreviewAgentPoolAddDecorator(
+            self.cmd,
+            self.client,
+            {"mode": CONST_NODEPOOL_MODE_MANAGEDSYSTEM},
+            self.resource_type,
+            self.agentpool_decorator_mode,
+        )
+
+        # fail on passing the wrong agentpool object
+        with self.assertRaises(CLIInternalError):
+            dec_1.set_up_managed_system_mode(None)
+
+        # Create an agentpool with various properties set
+        agentpool_1 = self.create_initialized_agentpool_instance(
+            restore_defaults=False,
+            count=3,
+            vm_size="Standard_D2s_v3",
+            os_type="Linux",
+            enable_auto_scaling=True,
+            min_count=1,
+            max_count=5,
+        )
+
+        # Store the original name for verification
+        original_name = agentpool_1.name
+
+        dec_1.context.attach_agentpool(agentpool_1)
+        dec_agentpool_1 = dec_1.set_up_managed_system_mode(agentpool_1)
+
+        # Verify that mode is set to ManagedSystem
+        self.assertEqual(dec_agentpool_1.mode, CONST_NODEPOOL_MODE_MANAGEDSYSTEM)
+
+        # 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}'")
+
+        # Test case 2: mode is not ManagedSystem - should return agentpool unchanged
+        dec_2 = AKSPreviewAgentPoolAddDecorator(
+            self.cmd,
+            self.client,
+            {"mode": CONST_NODEPOOL_MODE_SYSTEM},
+            self.resource_type,
+            self.agentpool_decorator_mode,
+        )
+
+        agentpool_2 = self.create_initialized_agentpool_instance(
+            restore_defaults=False,
+            count=3,
+            vm_size="Standard_D2s_v3",
+            mode=CONST_NODEPOOL_MODE_SYSTEM,
+        )
+
+        # Store reference to compare
+        original_agentpool_2 = agentpool_2
+
+        dec_2.context.attach_agentpool(agentpool_2)
+        dec_agentpool_2 = dec_2.set_up_managed_system_mode(agentpool_2)
+
+        # Verify that agentpool is returned unchanged
+        self.assertEqual(dec_agentpool_2, original_agentpool_2)
+
+        # Test case 3: mode is None (default) - should return agentpool unchanged
+        dec_3 = AKSPreviewAgentPoolAddDecorator(
+            self.cmd,
+            self.client,
+            {},  # No mode specified
+            self.resource_type,
+            self.agentpool_decorator_mode,
+        )
+
+        agentpool_3 = self.create_initialized_agentpool_instance(
+            restore_defaults=False,
+            count=3,
+            vm_size="Standard_D2s_v3",
+        )
+
+        original_agentpool_3 = agentpool_3
+
+        dec_3.context.attach_agentpool(agentpool_3)
+        dec_agentpool_3 = dec_3.set_up_managed_system_mode(agentpool_3)
+
+        # Verify that agentpool is returned unchanged
+        self.assertEqual(dec_agentpool_3, original_agentpool_3)
+
+    def common_construct_agentpool_profile_preview_with_managed_system_mode(self):
+        """Test that construct_agentpool_profile_preview properly handles ManagedSystem mode"""
+
+        # Test that when mode is ManagedSystem, only name and mode are preserved,
+        # and all other property setup methods are bypassed
+        dec = AKSPreviewAgentPoolAddDecorator(
+            self.cmd,
+            self.client,
+            {
+                "nodepool_name": "testnp",
+                "mode": CONST_NODEPOOL_MODE_MANAGEDSYSTEM,
+                # Add some parameters that would normally set properties
+                "node_count": 3,
+                "node_vm_size": "Standard_D2s_v3",
+                "enable_custom_ca_trust": True,
+                "crg_id": "test_crg_id",
+                "enable_artifact_streaming": True,
+            },
+            self.resource_type,
+            self.agentpool_decorator_mode,
+        )
+
+        # Construct the agentpool profile with mocked Azure API calls
+        with patch(
+            "azext_aks_preview.agentpool_decorator.cf_agent_pools",
+            return_value=Mock(list=Mock(return_value=[])),
+        ):
+            agentpool = dec.construct_agentpool_profile_preview()
+
+        # Verify that mode is set to ManagedSystem
+        self.assertEqual(agentpool.mode, CONST_NODEPOOL_MODE_MANAGEDSYSTEM)
+
+        # Verify that name is preserved
+        self.assertEqual(agentpool.name, "testnp")
+
+        # Verify that all other properties are None (bypassed)
+        for attr_name in vars(agentpool):
+            if attr_name not in ['name', 'mode'] and not attr_name.startswith('_'):
+                attr_value = getattr(agentpool, attr_name)
+                self.assertIsNone(attr_value,
+                    f"Attribute '{attr_name}' should be None but was '{attr_value}' when mode is ManagedSystem")
+
 
 class AKSPreviewAgentPoolAddDecoratorStandaloneModeTestCase(
     AKSPreviewAgentPoolAddDecoratorCommonTestCase
@@ -1399,6 +1611,9 @@ def test_set_up_agentpool_gateway_profile(self):
     def test_set_up_virtual_machines_profile(self):
         self.common_set_up_virtual_machines_profile()
 
+    def test_set_up_managed_system_mode(self):
+        self.common_set_up_managed_system_mode()
+
     def test_construct_agentpool_profile_preview(self):
         import inspect
 
@@ -1473,6 +1688,9 @@ def test_construct_agentpool_profile_preview(self):
 
         dec_1.context.raw_param.print_usage_statistics()
 
+    def test_construct_agentpool_profile_preview_with_managed_system_mode(self):
+        self.common_construct_agentpool_profile_preview_with_managed_system_mode()
+
 
 class AKSPreviewAgentPoolAddDecoratorManagedClusterModeTestCase(
     AKSPreviewAgentPoolAddDecoratorCommonTestCase
@@ -1518,10 +1736,13 @@ def test_set_up_agentpool_windows_profile(self):
 
     def test_set_up_agentpool_gateway_profile(self):
         self.common_set_up_agentpool_gateway_profile()
-    
+
     def test_set_up_virtual_machines_profile(self):
         self.common_set_up_virtual_machines_profile()
 
+    def test_set_up_managed_system_mode(self):
+        self.common_set_up_managed_system_mode()
+
     def test_construct_agentpool_profile_preview(self):
         import inspect
 
@@ -1885,7 +2106,7 @@ def common_update_fips_image(self):
         # fail on passing the wrong agentpool object
         with self.assertRaises(CLIInternalError):
             dec_2.update_fips_image(None)
-        
+
         agentpool_2 = self.create_initialized_agentpool_instance(enable_fips=True)
         dec_2.context.attach_agentpool(agentpool_2)
         dec_agentpool_2 = dec_2.update_fips_image(agentpool_2)