Merge pull request #100 from release-engineering/security_type

Azure: Fix bug with security_type on ARM64

Author: JAVGan

cloudpub/ms_azure/utils.py

@@ -281,6 +281,12 @@ def get_skuid(arch):
             return plan_name
         return f"{plan_name}-{arch.lower()}"
 
+    def get_safe_security_type(image_type):
+        # Arches which aren't x86Gen2 (like ARM64) doesn't work well with security type
+        if image_type != "x64Gen2":
+            return None
+        return security_type
+
     sku_mapping: Dict[str, str] = {}
     # Update the SKUs for each image in DiskVersions if needed
     for disk_version in disk_versions:
@@ -301,7 +307,7 @@ def get_skuid(arch):
 
     # Return the expected SKUs list
     res = [
-        VMISku.from_json({"image_type": k, "id": v, "security_type": security_type})
+        VMISku.from_json({"image_type": k, "id": v, "security_type": get_safe_security_type(k)})
         for k, v in sku_mapping.items()
     ]
     return sorted(res, key=attrgetter("id"))

tests/ms_azure/test_utils.py

@@ -379,7 +379,7 @@ def test_update_existing_skus_x86_gen1_default(
         assert res == [
             VMISku.from_json(x)
             for x in [
-                {"imageType": "x64Gen1", "skuId": "plan1", "securityType": ["trusted"]},
+                {"imageType": "x64Gen1", "skuId": "plan1"},
                 {"imageType": "x64Gen2", "skuId": "plan1-gen2", "securityType": ["trusted"]},
             ]
         ]
@@ -426,6 +426,40 @@ def test_update_existing_skus_x86_gen2_single(
             ]
         ]
 
+    def test_update_existing_skus_mixed_arches(
+        self, gen1_image: Dict[str, Any], gen2_image: Dict[str, Any], arm_image: Dict[str, Any]
+    ) -> None:
+        """Ensure the SKUs are properly made for disk versions using x86 and ARM arches."""
+        disk_versions = [
+            DiskVersion.from_json(
+                {
+                    "versionNumber": "2.0.0",
+                    "vmImages": [gen1_image, gen2_image, arm_image],
+                    "lifecycleState": "generallyAvailable",
+                }
+            )
+        ]
+        skus = [
+            VMISku.from_json(
+                {"imageType": "x64Gen2", "skuId": "plan1", "securityType": ["trusted"]}
+            )
+        ]
+        res = update_skus(
+            disk_versions=disk_versions,
+            generation="V2",
+            plan_name="plan1",
+            old_skus=skus,
+        )
+
+        assert res == [
+            VMISku.from_json(x)
+            for x in [
+                {"imageType": "x64Gen2", "skuId": "plan1", "securityType": ["trusted"]},
+                {"imageType": "arm64Gen2", "skuId": "plan1-arm64"},
+                {"imageType": "x64Gen1", "skuId": "plan1-gen1"},
+            ]
+        ]
+
     @pytest.mark.parametrize("generation", ["V1", "V2"])
     def test_update_existing_skus_arm64_single(
         self,