Merge pull request #56 from release-engineering/fixazure

Azure: Fix update_skus rountine for Gen1 default

Author: JAVGan

cloudpub/ms_azure/utils.py

@@ -235,6 +235,38 @@ def prepare_vm_images(
         return [VMImageDefinition.from_json(json_gen1)]
 
 
+def _build_skus(
+    disk_versions: List[DiskVersion],
+    default_gen: str,
+    alt_gen: str,
+    plan_name: str,
+    security_type: Optional[List[str]] = None,
+) -> List[VMISku]:
+    sku_mapping: Dict[str, str] = {}
+    # Update the SKUs for each image in DiskVersions if needed
+    for disk_version in disk_versions:
+        # Each disk version may have multiple images (Gen1 / Gen2)
+        for vmid in disk_version.vm_images:
+            # We'll name the main generation SKU as "{plan_name}" and
+            # the alternate generation SKU as "{plan-name}-genX"
+            arch = vmid.image_type.split("Gen")[0]
+            new_img_type = get_image_type_mapping(arch, default_gen)
+            new_img_alt_type = get_image_type_mapping(arch, alt_gen)
+
+            # we just want to add SKU whenever it's not set
+            if vmid.image_type == new_img_type:
+                sku_mapping.setdefault(new_img_type, plan_name)
+            elif vmid.image_type == new_img_alt_type:
+                sku_mapping.setdefault(new_img_alt_type, f"{plan_name}-gen{alt_gen[1:]}")
+
+    # Return the expected SKUs list
+    res = [
+        VMISku.from_json({"image_type": k, "id": v, "security_type": security_type})
+        for k, v in sku_mapping.items()
+    ]
+    return sorted(res, key=attrgetter("id"))
+
+
 def update_skus(
     disk_versions: List[DiskVersion],
     generation: str,
@@ -257,33 +289,33 @@ def update_skus(
     Returns:
         The updated list with VMISkus.
     """
-    sku_mapping: Dict[str, str] = {}
-    # All SKUs must have the same security_type thus picking the first one is OK
-    security_type = old_skus[0].security_type if old_skus else None
-
-    # Update the SKUs for each image in DiskVersions
-    for disk_version in disk_versions:
-        # Each disk version may have multiple images (Gen1 / Gen2)
-        for vmid in disk_version.vm_images:
-            # We'll name the main generation SKU as "{plan_name}" and
-            # the alternate generation SKU as "{plan-name}-genX"
-            alt_gen = 2 if generation == "V1" else 1
-            arch = vmid.image_type.split("Gen")[0]
-            new_img_type = get_image_type_mapping(arch, generation)
-            new_img_alt_type = get_image_type_mapping(arch, f"V{alt_gen}")
-
-            # we just want to add SKU whenever it's not set
-            if vmid.image_type == new_img_type:
-                sku_mapping.setdefault(new_img_type, plan_name)
-            elif vmid.image_type == new_img_alt_type:
-                sku_mapping.setdefault(new_img_alt_type, f"{plan_name}-gen{alt_gen}")
+    if not old_skus:
+        alt_gen = "V2" if generation == "V1" else "V1"
+        return _build_skus(
+            disk_versions, default_gen=generation, alt_gen=alt_gen, plan_name=plan_name
+        )
 
-    # Return the expected SKUs list
-    res = [
-        VMISku.from_json({"image_type": k, "id": v, "security_type": security_type})
-        for k, v in sku_mapping.items()
-    ]
-    return sorted(res, key=attrgetter("id"))
+    # The security type may exist only for Gen2, so it iterates over all gens to find it
+    security_type = None
+    # The alternate plan name ends with the suffix "-genX" and we can't change that once
+    # the offer is live, otherwise it will raise "BadRequest" with the message:
+    # "The property 'PlanId' is locked by a previous submission".
+    default_gen = "V2"
+    alt_gen = "V1"
+    for osku in old_skus:
+        if osku.security_type is not None:
+            security_type = osku.security_type
+        if osku.id.endswith("-gen2"):  # alternate is gen2 hence V1 is the default.
+            default_gen = "V1"
+            alt_gen = "V2"
+
+    return _build_skus(
+        disk_versions,
+        default_gen=default_gen,
+        alt_gen=alt_gen,
+        plan_name=plan_name,
+        security_type=security_type,
+    )
 
 
 def create_disk_version_from_scratch(

tests/ms_azure/test_utils.py

@@ -11,6 +11,7 @@
     VMImageDefinition,
     VMImageSource,
     VMIPlanTechConfig,
+    VMISku,
 )
 from cloudpub.ms_azure.utils import (
     AzurePublishingMetadata,
@@ -198,18 +199,77 @@ def test_prepare_vm_images_empty(
                 )
             assert expected_err in caplog.text
 
-    def test_update_skus(
+    def test_update_new_skus_gen2_default(
         self, technical_config_obj: VMIPlanTechConfig, metadata_azure_obj: AzurePublishingMetadata
     ) -> None:
         """Test for a "SKU update" from scratch."""
-        skus = technical_config_obj.skus
+        expected = technical_config_obj.skus
 
         res = update_skus(
             disk_versions=technical_config_obj.disk_versions,
             generation=metadata_azure_obj.generation,
             plan_name="plan-1",
         )
-        assert res == skus
+        assert res == expected
+
+    def test_update_new_skus_gen1_default(self, technical_config_obj: VMIPlanTechConfig) -> None:
+        """Test for a "SKU update" from scratch."""
+        expected = [
+            VMISku.from_json(x)
+            for x in [
+                {"imageType": "x64Gen1", "skuId": "plan1", "security_type": None},
+                {"imageType": "x64Gen2", "skuId": "plan1-gen2", "security_type": None},
+            ]
+        ]
+        res = update_skus(
+            disk_versions=technical_config_obj.disk_versions,
+            generation="V1",
+            plan_name="plan1",
+        )
+        assert res == expected
+
+    def test_update_existing_skus_gen2_default(
+        self, technical_config_obj: VMIPlanTechConfig, metadata_azure_obj: AzurePublishingMetadata
+    ) -> None:
+        """Test for a "SKU update" from scratch."""
+        res = update_skus(
+            disk_versions=technical_config_obj.disk_versions,
+            generation=metadata_azure_obj.generation,
+            plan_name="plan-1",
+            old_skus=technical_config_obj.skus,
+        )
+        assert res == [
+            VMISku.from_json(x)
+            for x in [
+                {"imageType": "x64Gen2", "skuId": "plan-1", "securityType": None},
+                {"imageType": "x64Gen1", "skuId": "plan-1-gen1", "securityType": None},
+            ]
+        ]
+
+    def test_update_existing_skus_gen1_default(
+        self, technical_config_obj: VMIPlanTechConfig
+    ) -> None:
+        skus = [
+            VMISku.from_json(x)
+            for x in [
+                {"imageType": "x64Gen1", "skuId": "plan1"},
+                {"imageType": "x64Gen2", "skuId": "plan1-gen2", "securityType": ["trusted"]},
+            ]
+        ]
+        technical_config_obj.skus = skus
+        res = update_skus(
+            disk_versions=technical_config_obj.disk_versions,
+            generation="V1",
+            plan_name="plan1",
+            old_skus=technical_config_obj.skus,
+        )
+        assert res == [
+            VMISku.from_json(x)
+            for x in [
+                {"imageType": "x64Gen1", "skuId": "plan1", "securityType": ["trusted"]},
+                {"imageType": "x64Gen2", "skuId": "plan1-gen2", "securityType": ["trusted"]},
+            ]
+        ]
 
     def test_create_disk_version_from_scratch(
         self,