uefi: mem: unit test for MemoryMapOwned + streamline test usage

Simplify MemoryMapOwned::from_initialized_mem to simplify and streamline the
construction in unit tests.

Author: phip1611

uefi/src/mem/memory_map/impl_.rs

@@ -102,10 +102,7 @@ pub struct MemoryMapRefMut<'a> {
 
 impl<'a> MemoryMapRefMut<'a> {
     /// Constructs a new [`MemoryMapRefMut`].
-    pub fn new(
-        buffer: &'a mut [u8],
-        meta: MemoryMapMeta,
-    ) -> Result<Self, MemoryMapError> {
+    pub fn new(buffer: &'a mut [u8], meta: MemoryMapMeta) -> Result<Self, MemoryMapError> {
         if buffer.as_ptr().align_offset(8) != 0 {
             return Err(MemoryMapError::Misaligned);
         }
@@ -301,10 +298,12 @@ impl MemoryMapBackingMemory {
         Self(slice)
     }
 
+    /// INTERNAL, for unit tests.
+    ///
     /// Creates an instance from the provided memory, which is not necessarily
     /// on the UEFI heap.
     #[cfg(test)]
-    fn from_slice(buffer: &mut [u8]) -> Self {
+    pub(crate) fn from_slice(buffer: &mut [u8]) -> Self {
         let len = buffer.len();
         unsafe { Self::from_raw(buffer.as_mut_ptr(), len) }
     }
@@ -347,6 +346,10 @@ impl Drop for MemoryMapBackingMemory {
                 log::error!("Failed to deallocate memory map: {e:?}");
             }
         } else {
+            #[cfg(test)]
+            log::debug!("Boot services are not available in unit tests.");
+
+            #[cfg(not(test))]
             log::debug!("Boot services are excited. Memory map won't be freed using the UEFI boot services allocator.");
         }
     }
@@ -362,30 +365,13 @@ pub struct MemoryMapOwned {
 }
 
 impl MemoryMapOwned {
-    /// Creates a [`MemoryMapOwned`] from the give initialized memory map behind
-    /// the buffer and the reported `desc_size` from UEFI.
+    /// Creates a [`MemoryMapOwned`] from the given **initialized** memory map
+    /// (stored inside the provided buffer) and the corresponding
+    /// [`MemoryMapMeta`].
     pub(crate) fn from_initialized_mem(buf: MemoryMapBackingMemory, meta: MemoryMapMeta) -> Self {
         assert!(meta.desc_size >= mem::size_of::<MemoryDescriptor>());
         let len = meta.entry_count();
-        MemoryMapOwned {
-            buf,
-            meta,
-            len,
-        }
-    }
-
-    #[cfg(test)]
-    pub(super) fn from_raw(buf: &mut [u8], desc_size: usize) -> Self {
-        let mem = MemoryMapBackingMemory::from_slice(buf);
-        Self::from_initialized_mem(
-            mem,
-            MemoryMapMeta {
-                map_size: buf.len(),
-                desc_size,
-                map_key: MemoryMapKey(0),
-                desc_version: MemoryDescriptor::VERSION,
-            },
-        )
+        MemoryMapOwned { buf, meta, len }
     }
 }
 
@@ -523,4 +509,22 @@ mod tests {
         mmap.sort();
         assert!(mmap.is_sorted());
     }
+
+    /// Basic sanity checks for the type [`MemoryMapOwned`].
+    #[test]
+    fn memory_map_owned() {
+        let mut memory = new_mmap_memory();
+        let (mmap, meta) = mmap_raw(&mut memory);
+        let mmap = MemoryMapBackingMemory::from_slice(mmap);
+        let mut mmap = MemoryMapOwned::from_initialized_mem(mmap, meta);
+
+        assert_eq!(mmap.entries().count(), 3);
+        assert_eq!(
+            mmap.entries().copied().collect::<Vec<_>>().as_slice(),
+            &BASE_MMAP_UNSORTED
+        );
+        assert!(!mmap.is_sorted());
+        mmap.sort();
+        assert!(mmap.is_sorted());
+    }
 }

uefi/src/mem/memory_map/mod.rs

@@ -106,13 +106,21 @@ mod tests_mmap_artificial {
     use core::mem::{size_of, size_of_val};
 
     fn buffer_to_map(buffer: &mut [MemoryDescriptor]) -> MemoryMapOwned {
-        let byte_buffer = {
-            unsafe {
-                core::slice::from_raw_parts_mut(buffer.as_mut_ptr() as *mut u8, size_of_val(buffer))
-            }
+        let mmap_len = size_of_val(buffer);
+        let mmap = {
+            unsafe { core::slice::from_raw_parts_mut(buffer.as_mut_ptr() as *mut u8, mmap_len) }
         };
 
-        MemoryMapOwned::from_raw(byte_buffer, size_of::<MemoryDescriptor>())
+        let mmap = MemoryMapBackingMemory::from_slice(mmap);
+        MemoryMapOwned::from_initialized_mem(
+            mmap,
+            MemoryMapMeta {
+                map_size: mmap_len,
+                desc_size: size_of::<MemoryDescriptor>(),
+                map_key: Default::default(),
+                desc_version: MemoryDescriptor::VERSION,
+            },
+        )
     }
 
     #[test]
@@ -259,7 +267,8 @@ mod tests_mmap_real {
         let mut buf = MMAP_RAW;
         let buf =
             unsafe { slice::from_raw_parts_mut(buf.as_mut_ptr().cast::<u8>(), MMAP_META.map_size) };
-        let mut mmap = MemoryMapOwned::from_raw(buf, MMAP_META.desc_size);
+        let buf = MemoryMapBackingMemory::from_slice(buf);
+        let mut mmap = MemoryMapOwned::from_initialized_mem(buf, MMAP_META);
         mmap.sort();
 
         let entries = mmap.entries().copied().collect::<Vec<_>>();