disk: handle when passed disktype is empty

microsoft · Mar 4, 2024 · 130127c · 130127c
1 parent b28f17b
commit 130127c
Show file tree

Hide file tree

Showing 5 changed files with 66 additions and 60 deletions.
diff --git a/lisa/features/availability.py b/lisa/features/availability.py
@@ -38,19 +38,15 @@ class AvailabilitySettings(schema.FeatureSettings):
             ],
         ),
         metadata=field_metadata(
-            decoder=lambda input: (
-                search_space.decode_set_space_by_type(
-                    data=input, base_type=AvailabilityType
-                )
-                if str(input).strip()
-                else search_space.SetSpace(
-                    is_allow_set=True,
-                    items=[
-                        AvailabilityType.NoRedundancy,
-                        AvailabilityType.AvailabilitySet,
-                        AvailabilityType.AvailabilityZone,
-                    ],
-                )
+            decoder=partial(
+                search_space.decode_nullable_set_space,
+                base_type=AvailabilityType,
+                default_values=[
+                    AvailabilityType.NoRedundancy,
+                    AvailabilityType.AvailabilitySet,
+                    AvailabilityType.AvailabilityZone,
+                ],
+                is_allow_set=True,
             )
         ),
     )
@@ -61,11 +57,12 @@ class AvailabilitySettings(schema.FeatureSettings):
             items=[],
         ),
         metadata=field_metadata(
-            decoder=lambda input: (
-                search_space.decode_set_space_by_type(data=input, base_type=int)
-                if str(input).strip()
-                else (search_space.SetSpace[int](is_allow_set=True, items=[]))
-            ),
+            decoder=partial(
+                search_space.decode_nullable_set_space,
+                base_type=int,
+                default_values=[],
+                is_allow_set=True,
+            )
         ),
     )
 

diff --git a/lisa/features/security_profile.py b/lisa/features/security_profile.py
@@ -46,19 +46,15 @@ class SecurityProfileSettings(schema.FeatureSettings):
             ],
         ),
         metadata=field_metadata(
-            decoder=lambda input: (
-                search_space.decode_set_space_by_type(
-                    data=input, base_type=SecurityProfileType
-                )
-                if str(input).strip()
-                else search_space.SetSpace(
-                    items=[
-                        SecurityProfileType.Standard,
-                        SecurityProfileType.SecureBoot,
-                        SecurityProfileType.CVM,
-                        SecurityProfileType.Stateless,
-                    ]
-                )
+            decoder=partial(
+                search_space.decode_nullable_set_space,
+                base_type=SecurityProfileType,
+                default_values=[
+                    SecurityProfileType.Standard,
+                    SecurityProfileType.SecureBoot,
+                    SecurityProfileType.CVM,
+                    SecurityProfileType.Stateless,
+                ],
             )
         ),
     )

diff --git a/lisa/schema.py b/lisa/schema.py
@@ -434,6 +434,21 @@ class DiskControllerType(str, Enum):
 ]
 
 
+os_disk_types: List[DiskType] = [
+    DiskType.StandardHDDLRS,
+    DiskType.StandardSSDLRS,
+    DiskType.Ephemeral,
+    DiskType.PremiumSSDLRS,
+]
+
+data_disk_types: List[DiskType] = [
+    DiskType.StandardHDDLRS,
+    DiskType.StandardSSDLRS,
+    DiskType.PremiumSSDLRS,
+    DiskType.UltraSSDLRS,
+]
+
+
 @dataclass_json()
 @dataclass()
 class DiskOptionSettings(FeatureSettings):
@@ -443,12 +458,7 @@ class DiskOptionSettings(FeatureSettings):
     ] = field(  # type:ignore
         default_factory=partial(
             search_space.SetSpace,
-            items=[
-                DiskType.StandardHDDLRS,
-                DiskType.StandardSSDLRS,
-                DiskType.Ephemeral,
-                DiskType.PremiumSSDLRS,
-            ],
+            items=os_disk_types,
         ),
         metadata=field_metadata(
             decoder=partial(search_space.decode_set_space_by_type, base_type=DiskType)
@@ -465,15 +475,14 @@ class DiskOptionSettings(FeatureSettings):
     ] = field(  # type:ignore
         default_factory=partial(
             search_space.SetSpace,
-            items=[
-                DiskType.StandardHDDLRS,
-                DiskType.StandardSSDLRS,
-                DiskType.PremiumSSDLRS,
-                DiskType.UltraSSDLRS,
-            ],
+            items=data_disk_types,
         ),
         metadata=field_metadata(
-            decoder=partial(search_space.decode_set_space_by_type, base_type=DiskType)
+            decoder=partial(
+                search_space.decode_nullable_set_space,
+                base_type=DiskType,
+                default_values=data_disk_types,
+            )
         ),
     )
     data_disk_count: search_space.CountSpace = field(
@@ -524,14 +533,10 @@ class DiskOptionSettings(FeatureSettings):
             items=[DiskControllerType.SCSI, DiskControllerType.NVME],
         ),
         metadata=field_metadata(
-            decoder=lambda input: (
-                search_space.decode_set_space_by_type(
-                    data=input, base_type=DiskControllerType
-                )
-                if str(input).strip()
-                else search_space.SetSpace(
-                    items=[DiskControllerType.SCSI, DiskControllerType.NVME]
-                )
+            decoder=partial(
+                search_space.decode_nullable_set_space,
+                base_type=DiskControllerType,
+                default_values=[DiskControllerType.SCSI, DiskControllerType.NVME],
             )
         ),
     )

diff --git a/lisa/search_space.py b/lisa/search_space.py
@@ -710,3 +710,15 @@ def create_set_space(
     else:
         set_space = None
     return set_space
+
+
+def decode_nullable_set_space(
+    data: Any, base_type: Any, default_values: Any, is_allow_set: bool = False
+) -> Any:
+    if str(data).strip():
+        return decode_set_space_by_type(data, base_type=base_type)
+    else:
+        return decode_set_space_by_type(
+            SetSpace(is_allow_set=is_allow_set, items=default_values),
+            base_type=base_type,
+        )
diff --git a/lisa/sut_orchestrator/azure/features.py b/lisa/sut_orchestrator/azure/features.py
@@ -2930,14 +2930,10 @@ class ArchitectureSettings(schema.FeatureSettings):
             search_space.SetSpace, items=[ArchitectureType.x64, ArchitectureType.Arm64]
         ),
         metadata=field_metadata(
-            decoder=lambda input: (
-                search_space.decode_set_space_by_type(
-                    data=input, base_type=ArchitectureType
-                )
-                if str(input).strip()
-                else search_space.SetSpace(
-                    items=[ArchitectureType.x64, ArchitectureType.Arm64]
-                )
+            decoder=partial(
+                search_space.decode_nullable_set_space,
+                base_type=ArchitectureType,
+                default_values=[ArchitectureType.x64, ArchitectureType.Arm64],
             )
         ),
     )