disks: Add guestimate of device partition paths

Signed-off-by: Ikey Doherty <[email protected]>

Author: ikeycode

crates/disks/src/lib.rs

@@ -61,6 +61,28 @@ impl BlockDevice {
         }
     }
 
+    /// Returns the path to the partition with the given index.
+    /// No attempt is made to verify the existence of the partition.
+    pub fn partition_path(&self, index: u32) -> PathBuf {
+        if let BlockDevice::Disk(disk) = self {
+            match **disk {
+                Disk::Scsi(_) | Disk::Virtual(_) => {
+                    // Add N to the end of the device name
+                    return disk.device_path().join(format!("{}{}", disk.name(), index));
+                }
+                Disk::Mock(ref d) if d.parts_prefix => {
+                    return PathBuf::from("/dev").join(format!("{}p{}", disk.name(), index));
+                }
+                Disk::Mock(_) => {
+                    return PathBuf::from("/dev").join(format!("{}{}", disk.name(), index));
+                }
+                _ => {}
+            }
+        }
+        // Add pN to the end of the device name
+        self.device().with_file_name(format!("{}p{}", self.name(), index))
+    }
+
     /// Creates a mock block device with a specified number of sectors.
     pub fn mock_device(disk: mock::MockDisk) -> Self {
         BlockDevice::Disk(Box::new(Disk::Mock(disk)))
@@ -181,4 +203,19 @@ mod tests {
             }
         }
     }
+
+    #[test]
+    fn test_partition_paths() {
+        // Create a mock SCSI disk
+        let scsi_disk = mock::MockDisk::new_with_name("sda", 1000, false);
+        let device = BlockDevice::mock_device(scsi_disk);
+        assert_eq!(device.partition_path(1).to_str().unwrap(), "/dev/sda1");
+        assert_eq!(device.partition_path(2).to_str().unwrap(), "/dev/sda2");
+
+        // Create a mock NVMe disk
+        let nvme_disk = mock::MockDisk::new_with_name("nvme0n1", 1000, true);
+        let device = BlockDevice::mock_device(nvme_disk);
+        assert_eq!(device.partition_path(1).to_str().unwrap(), "/dev/nvme0n1p1");
+        assert_eq!(device.partition_path(2).to_str().unwrap(), "/dev/nvme0n1p2");
+    }
 }

crates/disks/src/mock.rs

@@ -15,34 +15,45 @@ use crate::{partition::Partition, BasicDisk};
 ///
 /// This struct wraps a BasicDisk to provide mock functionality for testing.
 #[derive(Debug)]
-pub struct MockDisk(pub BasicDisk);
+pub struct MockDisk {
+    basic_disk: BasicDisk,
+    pub parts_prefix: bool,
+}
 
 impl Deref for MockDisk {
     type Target = BasicDisk;
 
     fn deref(&self) -> &Self::Target {
-        &self.0
+        &self.basic_disk
     }
 }
 
 impl MockDisk {
     /// Creates a new mock disk with the specified size in bytes
     pub fn new(size_bytes: u64) -> Self {
+        Self::new_with_name("mock0", size_bytes, false)
+    }
+
+    pub fn new_with_name(name: &str, size_bytes: u64, parts_prefix: bool) -> Self {
         let sectors = size_bytes / 512;
         let disk = BasicDisk {
-            name: "mock0".to_string(),
+            name: name.to_string(),
             sectors,
-            device: PathBuf::from("/dev/mock0"),
+            device: PathBuf::from(format!("/dev/{}", name)),
             model: Some("Mock Device".to_string()),
             vendor: Some("Mock Vendor".to_string()),
             partitions: Vec::new(),
         };
-        Self(disk)
+
+        Self {
+            basic_disk: disk,
+            parts_prefix,
+        }
     }
 
     /// Add a partition to the mock disk at the specified byte offsets
     pub fn add_partition(&mut self, start_bytes: u64, end_bytes: u64) {
-        let partition_number = self.0.partitions().len() + 1;
+        let partition_number = self.basic_disk.partitions().len() + 1;
         let start = start_bytes / 512;
         let end = end_bytes / 512;
 
@@ -56,6 +67,6 @@ impl MockDisk {
             device: PathBuf::from(format!("/dev/mock0p{}", partition_number)),
         };
 
-        self.0.partitions_mut().push(partition);
+        self.basic_disk.partitions_mut().push(partition);
     }
 }