Fixes to test cases and new test cases

src/aks-preview/azext_aks_preview/azurecontainerstorage/_consts.py

@@ -7,11 +7,11 @@
 CONST_ACSTOR_IO_ENGINE_LABEL_KEY = "acstor.azure.com/io-engine"
 CONST_ACSTOR_IO_ENGINE_LABEL_VAL = "acstor"
 CONST_ACSTOR_K8S_EXTENSION_NAME = "microsoft.azurecontainerstorage"
+CONST_DISK_TYPE_EPHEMERAL_VOLUME_ONLY = "EphemeralVolumeOnly"
+CONST_DISK_TYPE_PV_WITH_ANNOTATION = "PersistentVolumeWithAnnotation"
 CONST_EPHEMERAL_NVME_PERF_TIER_BASIC = "Basic"
 CONST_EPHEMERAL_NVME_PERF_TIER_PREMIUM = "Premium"
 CONST_EPHEMERAL_NVME_PERF_TIER_STANDARD = "Standard"
-CONST_DISK_TYPE_EPHEMERAL_VOLUME_ONLY = "EphemeralVolumeOnly"
-CONST_DISK_TYPE_PV_WITH_ANNOTATION = "PersistentVolumeWithAnnotation"
 CONST_EXT_INSTALLATION_NAME = "azurecontainerstorage"
 CONST_K8S_EXTENSION_CLIENT_FACTORY_MOD_NAME = "azext_k8s_extension._client_factory"
 CONST_K8S_EXTENSION_CUSTOM_MOD_NAME = "azext_k8s_extension.custom"

src/aks-preview/azext_aks_preview/azurecontainerstorage/_validators.py

@@ -410,7 +410,7 @@ def _validate_storage_pool_size(storage_pool_size, storage_pool_type):
             )
 
 
-def _validate_nodepools(
+def _validate_nodepools(  # pylint: disable=too-many-branches
     nodepool_list,
     agentpool_details,
     storage_pool_type,
@@ -431,12 +431,13 @@ def _validate_nodepools(
                 'in a cluster where Azure Container Storage is already installed.'
             )
 
-        for agentpool in agentpool_details:
-            node_labels = agentpool.get("node_labels")
-            if node_labels is not None and \
-               node_labels.get(CONST_ACSTOR_IO_ENGINE_LABEL_KEY) is not None:
-                nodepool_name = agentpool.get("name")
-                nodepool_arr.append(nodepool_name)
+        if agentpool_details is not None:
+            for agentpool in agentpool_details:
+                node_labels = agentpool.get("node_labels")
+                if node_labels is not None and \
+                   node_labels.get(CONST_ACSTOR_IO_ENGINE_LABEL_KEY) is not None:
+                    nodepool_name = agentpool.get("name")
+                    nodepool_arr.append(nodepool_name)
 
         if len(nodepool_arr) == 0:
             raise ArgumentUsageError(
@@ -453,8 +454,9 @@ def _validate_nodepools(
         )
     else:
         agentpool_names = []
-        for details in agentpool_details:
-            agentpool_names.append(details.get("name"))
+        if agentpool_details is not None:
+            for details in agentpool_details:
+                agentpool_names.append(details.get("name"))
         if not nodepool_list:
             agentpool_names_str = ', '.join(agentpool_names)
             raise RequiredArgumentMissingError(
@@ -480,7 +482,7 @@ def _validate_nodepools(
                     )
                 mode = agentpool.get("mode")
                 node_taints = agentpool.get("node_taints")
-                if mode.lower() == "system" and node_taints is not None:
+                if mode is not None and mode.lower() == "system" and node_taints is not None:
                     critical_taint = "CriticalAddonsOnly=true:NoSchedule"
                     if critical_taint.casefold() in (taint.casefold() for taint in node_taints):
                         raise InvalidArgumentValueError(
@@ -502,7 +504,8 @@ def _validate_nodepools(
                         nvme_nodepool_found = True
 
                 node_count = agentpool.get("count")
-                available_node_count = available_node_count + node_count
+                if node_count is not None:
+                    available_node_count = available_node_count + node_count
 
     if available_node_count < 3:
         raise UnknownError(

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

@@ -14,6 +14,7 @@
     InvalidArgumentValueError,
     MutuallyExclusiveArgumentError,
     RequiredArgumentMissingError,
+    UnknownError,
 )
 from azure.cli.core.util import CLIError
 
@@ -944,8 +945,7 @@ def test_enable_with_ephemeral_disk_nvme_perf_tier_and_non_ephemeral_disk_pool(s
         perf_tier = acstor_consts.CONST_EPHEMERAL_NVME_PERF_TIER_PREMIUM
         storage_pool_type = acstor_consts.CONST_STORAGE_POOL_TYPE_AZURE_DISK
         err = (
-            "Cannot set --ephemeral-disk-nvme-perf-tier when storage pool type: "
-            "ephemeralDisk option: NVMe is not enabled for Azure Container Storage."
+            "Cannot set --ephemeral-disk-nvme-perf-tier when --enable-azure-container-storage is not ephemeralDisk."
         )
         with self.assertRaises(ArgumentUsageError) as cm:
             acstor_validator.validate_enable_azure_container_storage_params(
@@ -959,8 +959,8 @@ def test_enable_with_ephemeral_disk_nvme_perf_tier_and_ephemeral_temp_disk_pool(
         storage_pool_type = acstor_consts.CONST_STORAGE_POOL_TYPE_EPHEMERAL_DISK
         storage_pool_option = acstor_consts.CONST_STORAGE_POOL_OPTION_SSD
         err = (
-            "Cannot set --ephemeral-disk-nvme-perf-tier when storage pool type: "
-            "ephemeralDisk option: NVMe is not enabled for Azure Container Storage. "
+            "Cannot set --ephemeral-disk-nvme-perf-tier along with --enable-azure-container-storage "
+            "when storage pool type: ephemeralDisk option: NVMe is not enabled for Azure Container Storage. "
             "Enable the option using --storage-pool-option."
         )
         with self.assertRaises(ArgumentUsageError) as cm:
@@ -1057,13 +1057,47 @@ def test_missing_nodepool_from_cluster_nodepool_list_multiple(self):
             )
         self.assertEqual(str(cm.exception), err)
 
+    def test_system_nodepool_with_taint(self):
+        storage_pool_name = "valid-name"
+        storage_pool_size = "5Ti"
+        storage_pool_type = acstor_consts.CONST_STORAGE_POOL_TYPE_EPHEMERAL_DISK
+        storage_pool_option = acstor_consts.CONST_STORAGE_POOL_OPTION_SSD
+        nodepool_list = "nodepool1"
+        agentpools = [{"name": "nodepool1", "mode": "System", "node_taints": ["CriticalAddonsOnly=true:NoSchedule"]}, {"name": "nodepool2", "count": 1}]
+        err = (
+            'Unable to install Azure Container Storage on system nodepool: nodepool1 '
+            'since it has a taint CriticalAddonsOnly=true:NoSchedule. Remove the taint from the node pool '
+            'and retry the Azure Container Storage operation.'
+        )
+        with self.assertRaises(InvalidArgumentValueError) as cm:
+            acstor_validator.validate_enable_azure_container_storage_params(
+                storage_pool_type, storage_pool_name, None, storage_pool_option, storage_pool_size, nodepool_list, agentpools, False, False, False, False, False, None, None, acstor_consts.CONST_DISK_TYPE_EPHEMERAL_VOLUME_ONLY, acstor_consts.CONST_EPHEMERAL_NVME_PERF_TIER_STANDARD
+            )
+        self.assertEqual(str(cm.exception), err)
+
+    def test_nodepool_from_cluster_nodepool_list_with_insufficient_count(self):
+        storage_pool_name = "valid-name"
+        storage_pool_size = "5Ti"
+        storage_pool_type = acstor_consts.CONST_STORAGE_POOL_TYPE_EPHEMERAL_DISK
+        storage_pool_option = acstor_consts.CONST_STORAGE_POOL_OPTION_SSD
+        nodepool_list = "nodepool1,nodepool2"
+        agentpools = [{"name": "nodepool1", "count": 1}, {"name": "nodepool2", "count": 1}]
+        err = (
+            "Insufficient nodes present. Azure Container Storage requires atleast 3 nodes to be enabled."
+        )
+        with self.assertRaises(UnknownError) as cm:
+            acstor_validator.validate_enable_azure_container_storage_params(
+                storage_pool_type, storage_pool_name, None, storage_pool_option, storage_pool_size, nodepool_list, agentpools, False, False, False, False, False, None, None, acstor_consts.CONST_DISK_TYPE_EPHEMERAL_VOLUME_ONLY, acstor_consts.CONST_EPHEMERAL_NVME_PERF_TIER_STANDARD
+            )
+        self.assertEqual(str(cm.exception), err)
+
     def test_valid_enable_for_azure_disk_pool(self):
         storage_pool_name = "valid-name"
         storage_pool_size = "5Ti"
         storage_pool_type = acstor_consts.CONST_STORAGE_POOL_TYPE_AZURE_DISK
         storage_pool_sku = acstor_consts.CONST_STORAGE_POOL_SKU_PREMIUM_LRS
         nodepool_list = "nodepool1,nodepool2"
-        agentpools = [{"name": "nodepool1"}, {"name": "nodepool2"}]
+        agentpools = [{"name": "nodepool1", "mode": "User", "count": 2}, {"name": "nodepool2", "mode": "System", "count": 1}]
         acstor_validator.validate_enable_azure_container_storage_params(
             storage_pool_type, storage_pool_name, storage_pool_sku, None, storage_pool_size, nodepool_list, agentpools, False, False, False, False, False, None, None, acstor_consts.CONST_DISK_TYPE_EPHEMERAL_VOLUME_ONLY, acstor_consts.CONST_EPHEMERAL_NVME_PERF_TIER_STANDARD
         )
@@ -1074,7 +1108,7 @@ def test_valid_enable_for_ephemeral_disk_pool(self):
         storage_pool_type = acstor_consts.CONST_STORAGE_POOL_TYPE_EPHEMERAL_DISK
         storage_pool_option = acstor_consts.CONST_STORAGE_POOL_OPTION_NVME
         nodepool_list = "nodepool1"
-        agentpools = [{"name": "nodepool1", "vm_size": "Standard_L8s_v3"}, {"name": "nodepool2", "vm_size": "Standard_L8s_v3"}]
+        agentpools = [{"name": "nodepool1", "vm_size": "Standard_L8s_v3", "mode": "System", "count": 5}, {"name": "nodepool2", "vm_size": "Standard_L8s_v3"}]
         acstor_validator.validate_enable_azure_container_storage_params(
             storage_pool_type, storage_pool_name, None, storage_pool_option, storage_pool_size, nodepool_list, agentpools, False, False, False, False, False, None, None, acstor_consts.CONST_DISK_TYPE_EPHEMERAL_VOLUME_ONLY, acstor_consts.CONST_EPHEMERAL_NVME_PERF_TIER_STANDARD
         )
@@ -1086,20 +1120,16 @@ def test_valid_enable_for_ephemeral_disk_pool_with_ephemeral_disk_volume_type(se
         storage_pool_option = acstor_consts.CONST_STORAGE_POOL_OPTION_NVME
         nodepool_list = "nodepool1"
         ephemeral_disk_volume_type = acstor_consts.CONST_DISK_TYPE_PV_WITH_ANNOTATION
-        agentpools = [{"name": "nodepool1", "vm_size": "Standard_L8s_v3"}, {"name": "nodepool2", "vm_size": "Standard_L8s_v3"}]
+        agentpools = [{"name": "nodepool1", "vm_size": "Standard_L8s_v3", "mode": "System", "count": 3}, {"name": "nodepool2", "vm_size": "Standard_L8s_v3"}]
         acstor_validator.validate_enable_azure_container_storage_params(
             storage_pool_type, storage_pool_name, None, storage_pool_option, storage_pool_size, nodepool_list, agentpools, False, False, False, False, False, ephemeral_disk_volume_type, None, acstor_consts.CONST_DISK_TYPE_EPHEMERAL_VOLUME_ONLY, acstor_consts.CONST_EPHEMERAL_NVME_PERF_TIER_STANDARD
         )
 
     def test_valid_enable_for_ephemeral_disk_pool_with_ephemeral_disk_volume_type_already_installed(self):
-        storage_pool_name = "valid-name"
-        storage_pool_size = "5Ti"
         storage_pool_type = acstor_consts.CONST_STORAGE_POOL_TYPE_EPHEMERAL_DISK
-        nodepool_list = "nodepool1"
         ephemeral_disk_volume_type = acstor_consts.CONST_DISK_TYPE_PV_WITH_ANNOTATION
-        agentpools = [{"name": "nodepool1", "vm_size": "Standard_L8s_v3"}, {"name": "nodepool2", "vm_size": "Standard_L8s_v3"}]
         acstor_validator.validate_enable_azure_container_storage_params(
-            storage_pool_type, None, None, None, None, None, None, True, False, False, False, False, ephemeral_disk_volume_type, None, acstor_consts.CONST_DISK_TYPE_EPHEMERAL_VOLUME_ONLY, acstor_consts.CONST_EPHEMERAL_NVME_PERF_TIER_STANDARD
+            storage_pool_type, None, None, None, None, None, None, True, False, False, True, False, ephemeral_disk_volume_type, None, acstor_consts.CONST_DISK_TYPE_EPHEMERAL_VOLUME_ONLY, acstor_consts.CONST_EPHEMERAL_NVME_PERF_TIER_STANDARD
         )
 
     def test_valid_enable_for_ephemeral_disk_pool_with_ephemeral_disk_nvme_perf_tier(self):
@@ -1109,18 +1139,14 @@ def test_valid_enable_for_ephemeral_disk_pool_with_ephemeral_disk_nvme_perf_tier
         storage_pool_option = acstor_consts.CONST_STORAGE_POOL_OPTION_NVME
         nodepool_list = "nodepool1"
         perf_tier = acstor_consts.CONST_EPHEMERAL_NVME_PERF_TIER_PREMIUM
-        agentpools = [{"name": "nodepool1", "vm_size": "Standard_L8s_v3"}, {"name": "nodepool2", "vm_size": "Standard_L8s_v3"}]
+        agentpools = [{"name": "nodepool1", "vm_size": "Standard_L8s_v3", "count": 4}, {"name": "nodepool2", "vm_size": "Standard_L8s_v3"}]
         acstor_validator.validate_enable_azure_container_storage_params(
             storage_pool_type, storage_pool_name, None, storage_pool_option, storage_pool_size, nodepool_list, agentpools, False, False, False, False, False, None, perf_tier, acstor_consts.CONST_DISK_TYPE_EPHEMERAL_VOLUME_ONLY, acstor_consts.CONST_EPHEMERAL_NVME_PERF_TIER_STANDARD
         )
 
     def test_valid_enable_for_ephemeral_disk_pool_with_azure_container_storage_per_tier_nvme_already_installed(self):
-        storage_pool_name = "valid-name"
-        storage_pool_size = "5Ti"
         storage_pool_type = acstor_consts.CONST_STORAGE_POOL_TYPE_EPHEMERAL_DISK
-        nodepool_list = "nodepool1"
         perf_tier = acstor_consts.CONST_EPHEMERAL_NVME_PERF_TIER_PREMIUM
-        agentpools = [{"name": "nodepool1", "vm_size": "Standard_L8s_v3"}, {"name": "nodepool2", "vm_size": "Standard_L8s_v3"}]
         acstor_validator.validate_enable_azure_container_storage_params(
             storage_pool_type, None, None, None, None, None, None, True, False, False, False, True, None, perf_tier, acstor_consts.CONST_DISK_TYPE_EPHEMERAL_VOLUME_ONLY, acstor_consts.CONST_EPHEMERAL_NVME_PERF_TIER_STANDARD
         )
@@ -1144,7 +1170,7 @@ def test_extension_installed_storagepool_type_installed(self):
         storage_pool_size = "5Ti"
         storage_pool_type = acstor_consts.CONST_STORAGE_POOL_TYPE_AZURE_DISK
         storage_pool_sku = acstor_consts.CONST_STORAGE_POOL_SKU_PREMIUM_LRS
-        agentpools = [{"name": "nodepool1", "node_labels": {"acstor.azure.com/io-engine": "acstor"}}, {"name": "nodepool2"}]
+        agentpools = [{"name": "nodepool1", "node_labels": {"acstor.azure.com/io-engine": "acstor"}, "count": 3}, {"name": "nodepool2"}]
         err = (
             "Invalid --enable-azure-container-storage value. "
             "Azure Container Storage is already enabled for storagepool type "
@@ -1161,7 +1187,7 @@ def test_valid_cluster_update(self):
         storage_pool_size = "5Ti"
         storage_pool_type = acstor_consts.CONST_STORAGE_POOL_TYPE_AZURE_DISK
         storage_pool_sku = acstor_consts.CONST_STORAGE_POOL_SKU_PREMIUM_LRS
-        agentpools = [{"name": "nodepool1", "node_labels": {"acstor.azure.com/io-engine": "acstor"}}, {"name": "nodepool2"}]
+        agentpools = [{"name": "nodepool1", "node_labels": {"acstor.azure.com/io-engine": "acstor"}, "mode": "User", "count": 3}, {"name": "nodepool2"}]
         acstor_validator.validate_enable_azure_container_storage_params(
             storage_pool_type, storage_pool_name, storage_pool_sku, None, storage_pool_size, None, agentpools, True, False, False, False, False, None, None, acstor_consts.CONST_DISK_TYPE_EPHEMERAL_VOLUME_ONLY, acstor_consts.CONST_EPHEMERAL_NVME_PERF_TIER_STANDARD
         )