1 files changed, 7 insertions, 15 deletions
diff --git a/kernel/src/test_support/page_mapper.cpp b/kernel/src/test_support/page_mapper.cpp
index abdcae5..9236deb 100644
--- a/kernel/src/test_support/page_mapper.cpp
+++ b/kernel/src/test_support/page_mapper.cpp
@@ -7,13 +7,12 @@
 #include <cstddef>
 #include <format>
 #include <stdexcept>
-#include <sys/mman.h>
 
 namespace kernel::tests
 {
 
-  page_mapper::page_mapper(kstd::units::bytes memory_size)
-      : memory{memory_size}
+  page_mapper::page_mapper(kstd::units::bytes physical_size, kstd::units::bytes virtual_size)
+      : memory{physical_size, virtual_size}
   {}
 
   auto page_mapper::map(kapi::memory::page page, kapi::memory::frame frame, flags) -> std::byte *
@@ -26,21 +25,14 @@ namespace kernel::tests
     }
 
     auto page_address = page.start_address();
-    auto sandbox_start = memory.heap_base();
-    auto sandbox_end = sandbox_start + memory.heap_size();
+    auto virtual_base = memory.virtual_base();
+    auto virtual_end = virtual_base + memory.virtual_size();
 
-    if (page_address >= sandbox_start && page_address < sandbox_end)
+    if (page_address >= virtual_base && page_address < virtual_end)
     {
       auto virtual_target = static_cast<std::byte *>(page_address);
       auto physical_offset = frame.number() * kapi::memory::frame::size;
-      auto mapped_ptr = mmap(virtual_target, kapi::memory::page::size.value, PROT_READ | PROT_WRITE,
-                             MAP_SHARED | MAP_FIXED, memory.memory_descriptor(), physical_offset.value);
-      if (mapped_ptr == MAP_FAILED)
-      {
-        throw std::runtime_error("Failed to map page");
-      }
-
-      return static_cast<std::byte *>(mapped_ptr);
+      return memory.map(kapi::memory::page::size, virtual_target, physical_offset.value);
     }
     else if (page_address >= kapi::memory::mmio_base)
     {
@@ -49,7 +41,7 @@ namespace kernel::tests
     else if (page_address >= kapi::memory::higher_half_direct_map_base)
     {
       auto offset = frame.number() * kapi::memory::frame::size;
-      return memory.ram_base() + offset;
+      return memory.physical_base() + offset;
     }
 
     return nullptr;