From 485bb9357bb9426a4bc95aeb381cb2d9527f1c39 Mon Sep 17 00:00:00 2001 From: 2Shirt <2xShirt@gmail.com> Date: Sun, 10 Nov 2024 21:54:08 -0800 Subject: [PATCH] Refactor add_partition_details --- src/system/diskpart.rs | 69 ++++++++++++++++++++------------------ src/tests/mod.rs | 66 +++++++++++++++++++----------------- src/tests/sample_output.rs | 8 +++-- 3 files changed, 77 insertions(+), 66 deletions(-) diff --git a/src/system/diskpart.rs b/src/system/diskpart.rs index a4bee11..ab4e397 100644 --- a/src/system/diskpart.rs +++ b/src/system/diskpart.rs @@ -71,22 +71,21 @@ pub fn get_disk_details(disk_id: &str, disk_size: u64, disk_details: Option<&str disk } -pub fn add_partition_details( - disk: &mut Disk, +pub fn get_partition_details( + disk_id: &str, disk_details: Option<&str>, part_details: Option<&str>, -) { - // TODO: Convert to get_partition_details(disk_id: &str, disk_details [..]){} - // Drops need to have mutable access to disk +) -> Vec { static RE_LIS: Lazy = Lazy::new(|| Regex::new(r"Partition\s+(\d+)\s+\w+\s+(\d+\s+\w+B)").unwrap()); + let mut parts = Vec::new(); // List partition let contents: String; if let Some(details) = disk_details { contents = String::from(details); } else { - let script = format!("select disk {}\r\nlist partition", disk.get_id()); + let script = format!("select disk {}\r\nlist partition", disk_id); contents = run_script(&script); }; for (_, [number, size]) in RE_LIS.captures_iter(&contents).map(|c| c.extract()) { @@ -95,15 +94,12 @@ pub fn add_partition_details( size: string_to_bytes(size), ..Default::default() }; - // pub fs_type: Option, - // pub label: Option, - // pub part_type: String, - disk.parts.push(part); + parts.push(part); } // Detail parititon - let mut script = vec![format!("select disk {}", disk.get_id())]; - for part in &disk.parts { + let mut script = vec![format!("select disk {}", disk_id)]; + for part in &parts { if part_details.is_some() { // Currently only used by tests break; @@ -120,7 +116,10 @@ pub fn add_partition_details( } else { part_contents = run_script(script.join("\r\n").as_str()); }; - parse_partition_details(disk, &part_contents); + parse_partition_details(&mut parts, &part_contents); + + // Done + parts } #[must_use] @@ -260,7 +259,7 @@ pub fn get_disks() -> Vec { if let Some(disk) = disks_map.remove(id) { // We remove the disk from the HashMap because we're moving it to the Vec let mut disk = get_disk_details(disk.id.as_str(), disk.size, Some(details)); - add_partition_details(&mut disk, Some(details), None); + disk.parts = get_partition_details(disk.get_id(), Some(details), None); disk.generate_descriptions(); disks_raw.push(disk); } @@ -286,7 +285,7 @@ pub fn parse_disk_numbers(contents: &str) -> Vec<&str> { disk_nums } -pub fn parse_partition_details(disk: &mut Disk, contents: &str) { +pub fn parse_partition_details(parts: &mut Vec, contents: &str) { // TODO: Update multiple fields at once? // https://stackoverflow.com/a/52905826 static RE_PAR: Lazy = Lazy::new(|| { @@ -307,25 +306,29 @@ pub fn parse_partition_details(disk: &mut Disk, contents: &str) { .map(|c| c.extract()) .enumerate() { - let part = &mut disk.parts[part_index]; - - // Partition info - if !part_type.trim().is_empty() { + if let Some(part) = parts.get_mut(part_index) { + // Partition info part.part_type = String::from(part_type.trim()); - } - // Volume info - for (_, [_id, letter, label, fs_type]) in - RE_VOL.captures_iter(vol_line).map(|c| c.extract()) - { - if !label.trim().is_empty() { - part.label = Some(String::from(label.trim())); - } - if !letter.trim().is_empty() { - part.letter = Some(String::from(letter.trim())); - } - if !fs_type.trim().is_empty() { - part.fs_type = Some(String::from(fs_type.trim())); + // Volume info + for (_, [_id, letter, label, fs_type]) in + RE_VOL.captures_iter(vol_line).map(|c| c.extract()) + { + part.label = if label.trim().is_empty() { + None + } else { + Some(String::from(label.trim())) + }; + part.letter = if letter.trim().is_empty() { + None + } else { + Some(String::from(letter.trim())) + }; + part.fs_type = if fs_type.trim().is_empty() { + None + } else { + Some(String::from(fs_type.trim())) + }; } } } @@ -335,7 +338,7 @@ pub fn refresh_disk_info(disk: &Disk) { // TODO: Needs refactor - assuming add_ functions are replaced with get_ variants info!("Refresh disk info"); let mut disk = get_disk_details(disk.get_id(), disk.size, None); - add_partition_details(&mut disk, None, None); + disk.parts = get_partition_details(disk.get_id(), None, None); disk.generate_descriptions(); } diff --git a/src/tests/mod.rs b/src/tests/mod.rs index 5b03d18..ab2fb29 100644 --- a/src/tests/mod.rs +++ b/src/tests/mod.rs @@ -13,7 +13,7 @@ // You should have received a copy of the GNU General Public License // along with Deja-vu. If not, see . // -mod sample_output; +pub mod sample_output; #[cfg(test)] mod diskpart { @@ -21,7 +21,7 @@ mod diskpart { use super::sample_output; #[test] - fn add_disk_details_gpt() { + fn get_disk_details_gpt() { let disk = system::diskpart::get_disk_details("", 0, Some(sample_output::DETAIL_DISK_GPT)); assert_eq!(disk.model, "Red Hat VirtIO SCSI Disk Device"); assert_eq!(disk.part_type, system::disk::PartitionTableType::Guid); @@ -29,7 +29,7 @@ mod diskpart { } #[test] - fn add_disk_details_mbr() { + fn get_disk_details_mbr() { let disk = system::diskpart::get_disk_details("", 0, Some(sample_output::DETAIL_DISK_MBR)); assert_eq!(disk.model, "Red Hat VirtIO SCSI Disk Device"); assert_eq!(disk.part_type, system::disk::PartitionTableType::Legacy); @@ -37,7 +37,34 @@ mod diskpart { } #[test] - fn add_partition_details() { + fn get_partition_details() { + // Left + let partition_1 = system::disk::Partition { + id: String::from("1"), + fs_type: Some(String::from("FAT32")), + label: Some(String::from("ESP")), + letter: None, + part_type: String::from("c12a7328-f81f-11d2-ba4b-00a0c93ec93b"), + size: 272629760, + ..Default::default() + }; + let partition_2 = system::disk::Partition { + id: String::from("2"), + part_type: String::from("e3c9e316-0b5c-4db8-817d-f92df00215ae"), + size: 16777216, + ..Default::default() + }; + let partition_4 = system::disk::Partition { + id: String::from("4"), + fs_type: Some(String::from("NTFS")), + label: Some(String::from("Windows")), + letter: Some(String::from("C")), + part_type: String::from("ebd0a0a2-b9e5-4433-87c0-68b6b72699c7"), + size: 50465865728, + ..Default::default() + }; + + // Right let mut disk = system::disk::Disk::default(); disk.parts.push(system::disk::Partition { id: String::from("1"), @@ -51,34 +78,13 @@ mod diskpart { id: String::from("4"), ..Default::default() }); - - let partition_1 = system::disk::Partition { - id: String::from("1"), - fs_type: Some(String::from("FAT32")), - label: Some(String::from("ESP")), - letter: None, - part_type: String::from("c12a7328-f81f-11d2-ba4b-00a0c93ec93b"), - ..Default::default() - }; - let partition_2 = system::disk::Partition { - id: String::from("2"), - part_type: String::from("e3c9e316-0b5c-4db8-817d-f92df00215ae"), - ..Default::default() - }; - let partition_4 = system::disk::Partition { - id: String::from("4"), - fs_type: Some(String::from("NTFS")), - label: Some(String::from("Windows")), - letter: Some(String::from("C")), - part_type: String::from("ebd0a0a2-b9e5-4433-87c0-68b6b72699c7"), - ..Default::default() - }; - - system::diskpart::add_partition_details( - &mut disk, + disk.parts = system::diskpart::get_partition_details( + "", Some(sample_output::DETAIL_DISK_GPT), Some(sample_output::SELECT_PART_DETAIL_PARTS), ); + + assert_eq!(disk.parts.len(), 3); assert_eq!(partition_1, disk.parts[0]); assert_eq!(partition_2, disk.parts[1]); assert_eq!(partition_4, disk.parts[2]); @@ -105,7 +111,7 @@ mod diskpart { }; system::diskpart::parse_partition_details( - &mut disk, + &mut disk.parts, &sample_output::SELECT_PART_DETAIL_ONE_PART, ); assert_eq!(partition_1, disk.parts[0]); diff --git a/src/tests/sample_output.rs b/src/tests/sample_output.rs index b554a3a..2d2681b 100644 --- a/src/tests/sample_output.rs +++ b/src/tests/sample_output.rs @@ -34,12 +34,14 @@ Clustered Disk : No Volume ### Ltr Label Fs Type Size Status Info ---------- --- ----------- ----- ---------- ------- --------- -------- - Volume 4 E Scratch NTFS Partition 59 GB Healthy + Volume 1 C Windows NTFS Partition 47 GB Healthy Boot + Volume 2 S ESP FAT32 Partition 260 MB Healthy System Partition ### Type Size Offset ------------- ---------------- ------- ------- - Partition 1 Reserved 16 MB 16 MB - Partition 2 Primary 59 GB 32 MB + Partition 1 Reserved 260 MB 1024 KB + Partition 2 Reserved 16 MB 261 MB + Partition 4 Primary 47 GB 277 MB Reached the end of the disk enumeration.