[Compute] Fix #23341: az vm list-skus: Fix filtering out VM sizes that are available regionally when they are restricted in all zones

src/azure-cli/azure/cli/command_modules/vm/_vm_utils.py

@@ -122,6 +122,37 @@ def _match_location(loc, locations):
     return result
 
 
+def is_sku_available(cmd, sku_info, zone):
+    """
+    The SKU is unavailable in the following cases:
+    1. regional restriction and the region is restricted
+    2. parameter --zone is input which indicates only showing skus with availability zones.
+       Meanwhile, zonal restriction and all zones are restricted
+    """
+    is_available = True
+    is_restrict_zone = False
+    is_restrict_location = False
+    if not sku_info.restrictions:
+        return is_available
+    for restriction in sku_info.restrictions:
+        if restriction.reason_code == 'NotAvailableForSubscription':
+            # The attribute location_info is not supported in versions 2017-03-30 and earlier
+            if cmd.supported_api_version(max_api='2017-03-30'):
+                is_available = False
+                break
+            if restriction.type == 'Zone' and not (
+                    set(sku_info.location_info[0].zones or []) - set(restriction.restriction_info.zones or [])):
+                is_restrict_zone = True
+            if restriction.type == 'Location' and (
+                    sku_info.location_info[0].location in (restriction.restriction_info.locations or [])):
+                is_restrict_location = True
+
+            if is_restrict_location or (is_restrict_zone and zone):
+                is_available = False
+                break
+    return is_available
+
+
 # pylint: disable=too-many-statements, too-many-branches
 def normalize_disk_info(image_data_disks=None,
                         data_disk_sizes_gb=None, attach_data_disks=None, storage_sku=None,

src/azure-cli/azure/cli/command_modules/vm/custom.py

@@ -1268,26 +1268,13 @@ def get_vm_details(cmd, resource_group_name, vm_name, include_user_data=False):
 
 
 def list_skus(cmd, location=None, size=None, zone=None, show_all=None, resource_type=None):
-    from ._vm_utils import list_sku_info
+    from ._vm_utils import list_sku_info, is_sku_available
     result = list_sku_info(cmd.cli_ctx, location)
     # pylint: disable=too-many-nested-blocks
     if not show_all:
         available_skus = []
         for sku_info in result:
-            is_available = True
-            if sku_info.restrictions:
-                for restriction in sku_info.restrictions:
-                    if restriction.reason_code == 'NotAvailableForSubscription':
-                        # The attribute location_info is not supported in versions 2017-03-30 and earlier
-                        if cmd.supported_api_version(max_api='2017-03-30'):
-                            is_available = False
-                            break
-                        # This SKU is not available only if all zones are restricted
-                        if not (set(sku_info.location_info[0].zones or []) -
-                                set(restriction.restriction_info.zones or [])):
-                            is_available = False
-                            break
-            if is_available:
+            if is_sku_available(cmd, sku_info, zone):
                 available_skus.append(sku_info)
         result = available_skus
     if resource_type:

src/azure-cli/azure/cli/command_modules/vm/tests/latest/test_vm_commands.py

@@ -1556,10 +1556,69 @@ def test_list_compute_skus_filter(self):
         result = self.cmd('vm list-skus -l westus --resource-type disks').get_output_in_json()
         self.assertTrue(result and len(result) == len([x for x in result if x['resourceType'] == 'disks']))
 
-    @AllowLargeResponse(size_kb=99999)
-    def test_list_compute_skus_partially_unavailable(self):
-        result = self.cmd('vm list-skus -l eastus --query "[?name==\'Standard_M64m\']"').get_output_in_json()
-        self.assertTrue(result and result[0]["restrictions"] and result[0]["restrictions"][0]["reasonCode"] == 'NotAvailableForSubscription')
+    @mock.patch('azure.cli.command_modules.vm._validators._compute_client_factory', autospec=True)
+    def test_list_compute_skus_partially_unavailable(self, client_factory_mock):
+        from azure.cli.core.mock import DummyCli
+        from azure.cli.command_modules.vm._vm_utils import is_sku_available
+        compute_client = mock.MagicMock()
+
+        cmd = mock.Mock()
+        cmd.supported_api_version = mock.Mock(return_value=False)
+        cmd.cli_ctx = DummyCli()
+
+        client_factory_mock.return_value = compute_client
+
+        np = mock.MagicMock()
+        location_info0 = mock.MagicMock()
+        location_info0.zones = [1, 2, 3]
+        location_info0.location = 'location'
+        location_info = [location_info0]
+        np.location_info = location_info
+        restriction_info_zone = mock.MagicMock()
+        restriction_info_zone.zones = [1, 2, 3]
+        restriction_zone = mock.MagicMock()
+        restriction_zone.reason_code = 'NotAvailableForSubscription'
+        restriction_zone.type = 'Zone'
+        restriction_zone.restriction_info = restriction_info_zone
+        restriction_info_location = mock.MagicMock()
+        restriction_info_location.locations = ['location']
+        restriction_location = mock.MagicMock()
+        restriction_location.reason_code = 'NotAvailableForSubscription'
+        restriction_location.type = 'Location'
+        restriction_location.restriction_info = restriction_info_location
+
+        # zonal restriction but not regional restriction
+        restrictions = [restriction_zone]
+        np.restrictions = restrictions
+        # show skus supporting availability zones
+        is_available = is_sku_available(cmd, np, True)
+        self.assertFalse(is_available)
+        # not show skus supporting availability zones
+        is_available = is_sku_available(cmd, np, None)
+        self.assertTrue(is_available)
+
+        # not all zones are restricted
+        restriction_info_zone.zones = [1, 2]
+        is_available = is_sku_available(cmd, np, True)
+        self.assertTrue(is_available)
+        is_available = is_sku_available(cmd, np, None)
+        self.assertTrue(is_available)
+
+        # regional restriction but not zonal restriction
+        restrictions = [restriction_location]
+        np.restrictions = restrictions
+        is_available = is_sku_available(cmd, np, True)
+        self.assertFalse(is_available)
+        is_available = is_sku_available(cmd, np, None)
+        self.assertFalse(is_available)
+
+        # regional restriction and zonal restriction
+        restrictions = [restriction_location, restriction_zone]
+        np.restrictions = restrictions
+        is_available = is_sku_available(cmd, np, True)
+        self.assertFalse(is_available)
+        is_available = is_sku_available(cmd, np, None)
+        self.assertFalse(is_available)
 
 
 class VMExtensionScenarioTest(ScenarioTest):