Rework unsupported security types

Introduce an environment-driven configuration for arches that don't support security type, expose it via a helper function, refactor security type handling to use this configuration, and add comprehensive tests for the new helper.

New Features:

Add UNSUPPORTED_SECURITY_TYPE_ARCHES constant to configure arches that don't support security type via environment variable
Introduce get_unsupported_security_type_arches function to parse the configured arches
Enhancements:

Update get_safe_security_type logic to skip unsupported arches based on the new function
Tests:

Add unit tests covering default, multiple, three-value, and empty cases for get_unsupported_security_type_arches

Signed-off-by: Jonathan Gangi <jgangi@redhat.com>

Assisted-by: Cursor/Gemini

Author: JAVGan

cloudpub/ms_azure/utils.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 import logging
+import os
 from operator import attrgetter
 from typing import Any, Dict, List, Optional, Tuple, TypedDict
 
@@ -21,6 +22,19 @@
 log = logging.getLogger(__name__)
 
 
+UNSUPPORTED_SECURITY_TYPE_ARCHES = (
+    os.environ.get("UNSUPPORTED_SECURITY_TYPE_ARCHES", "") or "x64Gen1"
+)
+"""
+The list of arches that don't support security type.
+
+This is a comma-separated list of arches that don't support security type.
+It's used to skip the security type for the arches that don't support it.
+It's set as an environment variable to allow for customization.
+If the environment variable is empty, the default "x64Gen1" is used.
+"""
+
+
 class AzurePublishingMetadata(PublishingMetadata):
     """A collection of metadata necessary for publishing a VHD Image into a product."""
 
@@ -303,6 +317,12 @@ def _all_skus_present(old_skus: List[VMISku], disk_versions: List[DiskVersion])
     return True
 
 
+def get_unsupported_security_type_arches() -> List[str]:
+    """Return the list of arches that don't support security type."""
+    arches = UNSUPPORTED_SECURITY_TYPE_ARCHES or "x64Gen1"
+    return [arch.strip() for arch in arches.split(",")]
+
+
 def _build_skus(
     disk_versions: List[DiskVersion],
     default_gen: str,
@@ -316,8 +336,8 @@ def get_skuid(arch):
         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":
+        # Some arches (like x86 Gen1) doesn't support security type, so we need to skip them.
+        if image_type in get_unsupported_security_type_arches():
             return None
         return security_type
 
@@ -348,8 +368,8 @@ def get_safe_security_type(image_type):
 
 
 def _get_security_type(old_skus: List[VMISku]) -> Optional[List[str]]:
-    # The security type may exist only for x64 Gen2, so it iterates over all gens to find it
-    # Get the security type for all gens
+    # The security type may not be applied for certain arches, like x64 Gen1.
+    # This function will return the proper security type for the arches that has it set.
     for osku in old_skus:
         if osku.security_type is not None:
             return osku.security_type

tests/ms_azure/test_utils.py

@@ -1,6 +1,7 @@
 import logging
 from operator import attrgetter
 from typing import Any, Dict
+from unittest import mock
 
 import pytest
 from _pytest.logging import LogCaptureFixture
@@ -17,6 +18,7 @@
     AzurePublishingMetadata,
     create_disk_version_from_scratch,
     get_image_type_mapping,
+    get_unsupported_security_type_arches,
     is_azure_job_not_complete,
     is_sas_present,
     prepare_vm_images,
@@ -489,7 +491,7 @@ def test_update_existing_skus_mixed_arches(
             VMISku.from_json(x)
             for x in [
                 {"imageType": "x64Gen2", "skuId": "plan1", "securityType": ["trusted"]},
-                {"imageType": "arm64Gen2", "skuId": "plan1-arm64"},
+                {"imageType": "arm64Gen2", "skuId": "plan1-arm64", "securityType": ["trusted"]},
                 {"imageType": "x64Gen1", "skuId": "plan1-gen1"},
             ]
         ]
@@ -580,3 +582,37 @@ def test_create_disk_version_from_scratch_arm64(
         res.vm_images = sorted(res.vm_images, key=attrgetter("image_type"))
 
         assert res == disk_version_arm64_obj
+
+    @mock.patch("cloudpub.ms_azure.utils.UNSUPPORTED_SECURITY_TYPE_ARCHES", "x64Gen1")
+    def test_get_unsupported_security_type_arches_default(self) -> None:
+        """Test that the function returns the default value when env var is not set."""
+        res = get_unsupported_security_type_arches()
+        assert res == ["x64Gen1"]
+
+    @mock.patch("cloudpub.ms_azure.utils.UNSUPPORTED_SECURITY_TYPE_ARCHES", "x64Gen1,arm64Gen1")
+    def test_get_unsupported_security_type_arches_multiple(self) -> None:
+        """Test that the function correctly splits comma-separated values."""
+        res = get_unsupported_security_type_arches()
+        assert res == ["x64Gen1", "arm64Gen1"]
+
+    @mock.patch(
+        "cloudpub.ms_azure.utils.UNSUPPORTED_SECURITY_TYPE_ARCHES", "x64Gen1,x64Gen2,arm64Gen1"
+    )
+    def test_get_unsupported_security_type_arches_multiple_three(self) -> None:
+        """Test that the function correctly handles three comma-separated values."""
+        res = get_unsupported_security_type_arches()
+        assert res == ["x64Gen1", "x64Gen2", "arm64Gen1"]
+
+    @mock.patch("cloudpub.ms_azure.utils.UNSUPPORTED_SECURITY_TYPE_ARCHES", "")
+    def test_get_unsupported_security_type_arches_empty(self) -> None:
+        """Test that the function returns default value when env var is empty."""
+        res = get_unsupported_security_type_arches()
+        assert res == ["x64Gen1"]
+
+    @mock.patch(
+        "cloudpub.ms_azure.utils.UNSUPPORTED_SECURITY_TYPE_ARCHES", "x64Gen1 , arm64Gen2 , x64Gen2"
+    )
+    def test_get_unsupported_security_type_arches_with_spaces(self) -> None:
+        """Test that the function correctly strips whitespace from comma-separated values."""
+        res = get_unsupported_security_type_arches()
+        assert res == ["x64Gen1", "arm64Gen2", "x64Gen2"]