Added "get_disc_type"

microsoft · Nov 22, 2024 · f423499 · f423499
1 parent 0756836
commit f423499
Show file tree

Hide file tree

Showing 6 changed files with 44 additions and 41 deletions.
diff --git a/lisa/features/disks.py b/lisa/features/disks.py
@@ -10,7 +10,7 @@
 from lisa.operating_system import BSD
 from lisa.tools import Mount
 from lisa.tools.mount import PartitionInfo
-from lisa.util import LisaException, get_matched_str
+from lisa.util import LisaException, constants, get_matched_str
 
 
 class Disk(Feature):
@@ -75,8 +75,8 @@ def get_resource_disk_mount_point(self) -> str:
     def get_resource_disks(self) -> List[str]:
         return []
 
-    def get_resource_disk_type(self) -> str:
-        return ""
+    def get_resource_disk_type(self) -> schema.ResourceDiskType:
+        return schema.ResourceDiskType.SCSI
 
     def get_luns(self) -> Dict[str, int]:
         raise NotImplementedError
@@ -105,25 +105,29 @@ def get_os_boot_partition(self) -> Optional[PartitionInfo]:
                 break
         return boot_partition
 
-    # Get disk controller type from the VM by checking the boot partition
-    def get_os_disk_controller_type(self) -> schema.DiskControllerType:
-        boot_partition = self.get_os_boot_partition()
-        assert boot_partition, "'boot_partition' must not be 'None'"
+    def get_disc_type(self, disk: str) -> schema.StorageInterfaceType:
         if isinstance(self._node.os, BSD):
-            if boot_partition.disk.startswith("da"):
-                os_disk_controller_type = schema.DiskControllerType.SCSI
-            elif boot_partition.disk.startswith("nvd"):
-                os_disk_controller_type = schema.DiskControllerType.NVME
+            if disk.startswith("da"):
+                disk_type = schema.StorageInterfaceType.SCSI
+            elif disk.startswith("nvd"):
+                disk_type = schema.StorageInterfaceType.NVME
             else:
-                raise LisaException(f"Unknown OS boot disk type {boot_partition.disk}")
+                raise LisaException(f"Unknown disk type {disk}")
         else:
-            if boot_partition.disk.startswith("nvme"):
-                os_disk_controller_type = schema.DiskControllerType.NVME
-            elif boot_partition.disk.startswith("sd"):
-                os_disk_controller_type = schema.DiskControllerType.SCSI
+            if disk.startswith("nvme"):
+                disk_type = schema.StorageInterfaceType.NVME
+            elif disk.startswith("sd"):
+                disk_type = schema.StorageInterfaceType.SCSI
             else:
-                raise LisaException(f"Unknown OS boot disk type {boot_partition.disk}")
-        return os_disk_controller_type
+                raise LisaException(f"Unknown disk type {disk}")
+        return disk_type
+
+    # Get disk controller type from the VM by checking the boot partition
+    def get_os_disk_controller_type(self) -> schema.DiskControllerType:
+        boot_partition = self.get_os_boot_partition()
+        assert boot_partition, "'boot_partition' must not be 'None'"
+        os_disk_controller_type = self.get_disc_type(boot_partition.disk)
+        return schema.DiskControllerType(os_disk_controller_type)
 
 
 DiskEphemeral = partial(

diff --git a/lisa/schema.py b/lisa/schema.py
@@ -431,9 +431,19 @@ class DiskType(str, Enum):
 ]
 
 
-class DiskControllerType(str, Enum):
-    SCSI = "SCSI"
+class StorageInterfaceType(str, Enum):
     NVME = "NVMe"
+    SCSI = "SCSI"
+
+
+class DiskControllerType(str, Enum):
+    SCSI = StorageInterfaceType.SCSI
+    NVME = StorageInterfaceType.NVME
+
+
+class ResourceDiskType(str, Enum):
+    SCSI = StorageInterfaceType.SCSI
+    NVME = StorageInterfaceType.NVME
 
 
 disk_controller_type_priority: List[DiskControllerType] = [

diff --git a/lisa/sut_orchestrator/azure/features.py b/lisa/sut_orchestrator/azure/features.py
@@ -1913,13 +1913,13 @@ def check_resource_disk_mounted(self) -> bool:
     # get resource disk type
     # function returns the type of resource disk/disks available on the VM
     # raises exception if no resource disk is available
-    def get_resource_disk_type(self) -> str:
+    def get_resource_disk_type(self) -> schema.ResourceDiskType:
         resource_disks = self.get_resource_disks()
         if not resource_disks:
             raise LisaException("No Resource disks are available on VM")
-        if constants.RESOURCE_DISK_TYPE_NVME in resource_disks[0]:
-            return constants.RESOURCE_DISK_TYPE_NVME
-        return constants.RESOURCE_DISK_TYPE_SCSI
+        return schema.ResourceDiskType(
+            self._node.features[Disk].get_disc_type(disk=resource_disks[0])
+        )
 
     def get_resource_disks(self) -> List[str]:
         resource_disk_list = []

diff --git a/lisa/util/constants.py b/lisa/util/constants.py
@@ -178,7 +178,3 @@
 SIGINT = 2
 SIGTERM = 15
 SIGKILL = 9
-
-# Resource Disk Types
-RESOURCE_DISK_TYPE_NVME = "nvme"
-RESOURCE_DISK_TYPE_SCSI = "scsi"
diff --git a/microsoft/testsuites/core/azure_image_standard.py b/microsoft/testsuites/core/azure_image_standard.py
@@ -13,6 +13,7 @@
     TestCaseMetadata,
     TestSuite,
     TestSuiteMetadata,
+    schema,
     simple_requirement,
 )
 from lisa.features import Disk
@@ -1116,10 +1117,7 @@ def verify_no_pre_exist_users(self, node: Node) -> None:
         ),
     )
     def verify_resource_disk_readme_file(self, node: RemoteNode) -> None:
-        if (
-            constants.RESOURCE_DISK_TYPE_NVME
-            == node.features[Disk].get_resource_disk_type()
-        ):
+        if schema.ResourceDiskType.NVME == node.features[Disk].get_resource_disk_type():
             raise SkippedException(
                 "Resource disk type is NVMe. NVMe disks are not formatted or mounted by"
                 " default and readme file wont be available"
@@ -1169,7 +1167,7 @@ def verify_resource_disk_readme_file(self, node: RemoteNode) -> None:
     )
     def verify_resource_disk_file_system(self, node: RemoteNode) -> None:
         node_disc = node.features[Disk]
-        if constants.RESOURCE_DISK_TYPE_NVME == node_disc.get_resource_disk_type():
+        if schema.ResourceDiskType.NVME == node.features[Disk].get_resource_disk_type():
             raise SkippedException(
                 "Resource disk type is NVMe. NVMe disks are not formatted or mounted by default"  # noqa: E501
             )

diff --git a/microsoft/testsuites/core/storage.py b/microsoft/testsuites/core/storage.py
@@ -14,6 +14,7 @@
     TestCaseMetadata,
     TestSuite,
     TestSuiteMetadata,
+    schema,
     simple_requirement,
 )
 from lisa.base_tools.service import Systemctl
@@ -161,10 +162,7 @@ def verify_disks_device_timeout_setting(
         ),
     )
     def verify_resource_disk_mounted(self, node: RemoteNode) -> None:
-        if (
-            constants.RESOURCE_DISK_TYPE_NVME
-            == node.features[Disk].get_resource_disk_type()
-        ):
+        if schema.ResourceDiskType.NVME == node.features[Disk].get_resource_disk_type():
             raise SkippedException(
                 "Resource disk type is NVMe. NVMe disks are not mounted by default"
             )
@@ -240,10 +238,7 @@ def verify_swap(self, node: RemoteNode) -> None:
         ),
     )
     def verify_resource_disk_io(self, node: RemoteNode) -> None:
-        if (
-            constants.RESOURCE_DISK_TYPE_NVME
-            == node.features[Disk].get_resource_disk_type()
-        ):
+        if schema.ResourceDiskType.NVME == node.features[Disk].get_resource_disk_type():
             raise SkippedException(
                 "Resource disk type is NVMe. NVMe has 'verify_nvme_function' and "
                 "'verify_nvme_function_unpartitioned' testcases to validate IO operations."  # noqa: E501