diff --git a/src/BUILD.gn b/src/BUILD.gn
index 1fb87d0..b39261d 100644
--- a/src/BUILD.gn
+++ b/src/BUILD.gn
@@ -1,13 +1,11 @@
 # Hypervisor specific code.
 source_set("src") {
   sources = [
-    "alloc.c",
     "api.c",
     "cpio.c",
     "cpu.c",
     "load.c",
     "main.c",
-    "mm.c",
     "vm.c",
   ]
 
@@ -16,6 +14,7 @@
     ":fdt",
     ":fdt_handler",
     ":memiter",
+    ":mm",
   ]
 
   if (is_debug) {
@@ -71,6 +70,17 @@
   ]
 }
 
+source_set("mm") {
+  sources = [
+    "alloc.c",
+    "mm.c",
+  ]
+
+  deps = [
+    ":common_debug",
+  ]
+}
+
 executable("unit_tests") {
   testonly = true
 
@@ -79,11 +89,24 @@
     "fdt_test.cc",
   ]
   cflags_cc = [ "-Wno-c99-extensions" ]
-
+  ldflags = [
+    "-Xlinker",
+    "-defsym=text_begin=0",
+    "-Xlinker",
+    "-defsym=text_end=100",
+    "-Xlinker",
+    "-defsym=rodata_begin=200",
+    "-Xlinker",
+    "-defsym=rodata_end=300",
+    "-Xlinker",
+    "-defsym=data_begin=400",
+    "-Xlinker",
+    "-defsym=data_end=500",
+  ]
   deps = [
     ":fdt",
     ":fdt_handler",
-    "mock:mm",
+    ":mm",
     "//third_party:gtest_main",
   ]
 }
diff --git a/src/arch/mock/inc/hf/arch/mm.h b/src/arch/mock/inc/hf/arch/mm.h
index 676b53f..3c016d0 100644
--- a/src/arch/mock/inc/hf/arch/mm.h
+++ b/src/arch/mock/inc/hf/arch/mm.h
@@ -5,10 +5,24 @@
 
 #include "hf/addr.h"
 
+/*
+ * Our mock architecture has page tables rather similar to aarch64, but not
+ * quite.
+ * - The highest level table is always 2, lowest level is 0.
+ * - Blocks are allowed at all levels.
+ * There are four types of entries:
+ * - Absent: 0
+ * - Page, at level 0: <page-aligned address> | <attrs> | 0x3
+ * - Block, at level 2 or 1: <block-aligned address> | <attrs> | 0x1
+ * - Subtable, at level 2 or 1: <subtable address> | 0x3
+ * <attrs> are always 0 for now.
+ */
+
 /* A page table entry. */
 typedef uint64_t pte_t;
 
 #define PAGE_LEVEL_BITS 9
+#define PL011_BASE 0x11
 
 /**
  * Returns the encoding of a page table entry that isn't present.
@@ -38,7 +52,7 @@
  */
 static inline pte_t arch_mm_block_pte(int level, paddr_t pa, uint64_t attrs)
 {
-	pte_t pte = pa_addr(pa) | attrs;
+	pte_t pte = pa_addr(pa) | attrs | 0x1;
 	if (level == 0) {
 		pte |= 0x2;
 	}
@@ -81,13 +95,20 @@
 	       (pte & 0x3) == (level == 0 ? 0x3 : 0x1);
 }
 
+static inline uint64_t hf_arch_mock_mm_clear_pte_attrs(pte_t pte)
+{
+	return pte & ~0x3;
+}
+
 /**
  * Clears the given physical address, i.e., sets the ignored bits (from a page
  * table perspective) to zero.
  */
 static inline paddr_t arch_mm_clear_pa(paddr_t pa)
 {
-	return pa_init(0);
+	/* This is assumed to round down to the page boundary. */
+	return pa_init(hf_arch_mock_mm_clear_pte_attrs(pa_addr(pa)) &
+		       ~((1 << PAGE_BITS) - 1));
 }
 
 /**
@@ -96,7 +117,7 @@
  */
 static inline paddr_t arch_mm_block_from_pte(pte_t pte)
 {
-	return pa_init(pte);
+	return pa_init(hf_arch_mock_mm_clear_pte_attrs(pte));
 }
 
 /**
@@ -105,7 +126,7 @@
  */
 static inline paddr_t arch_mm_table_from_pte(pte_t pte)
 {
-	return pa_init(pte);
+	return pa_init(hf_arch_mock_mm_clear_pte_attrs(pte));
 }
 
 /**
@@ -118,18 +139,53 @@
 }
 
 /**
+ * Given the attrs from a table at some level and the attrs from all the blocks
+ * in that table, return equivalent attrs to use for a block which will replace
+ * the entire table.
+ */
+static inline uint64_t arch_mm_combine_table_entry_attrs(uint64_t table_attrs,
+							 uint64_t block_attrs)
+{
+	return table_attrs | block_attrs;
+}
+
+/**
  * Invalidates stage-1 TLB entries referring to the given virtual address range.
  */
-void arch_mm_invalidate_stage1_range(vaddr_t va_begin, vaddr_t va_end);
+static inline void arch_mm_invalidate_stage1_range(vaddr_t va_begin,
+						   vaddr_t va_end)
+{
+}
 
 /**
  * Invalidates stage-2 TLB entries referring to the given intermediate physical
  * address range.
  */
-void arch_mm_invalidate_stage2_range(ipaddr_t va_begin, ipaddr_t va_end);
+static inline void arch_mm_invalidate_stage2_range(ipaddr_t va_begin,
+						   ipaddr_t va_end)
+{
+}
+
+/**
+ * Determines the maximum level supported by the given mode.
+ */
+static inline int arch_mm_max_level(int mode)
+{
+	(void)mode;
+	return 2;
+}
+
+static inline uint64_t arch_mm_mode_to_attrs(int mode)
+{
+	(void)mode;
+	return 0;
+}
+
+static inline bool arch_mm_init(paddr_t table, bool first)
+{
+	(void)table;
+	(void)first;
+	return true;
+}
 
 void arch_mm_set_vm(uint64_t vmid, paddr_t table);
-
-uint64_t arch_mm_mode_to_attrs(int mode);
-bool arch_mm_init(paddr_t table, bool first);
-int arch_mm_max_level(int mode);
diff --git a/src/fdt_handler_test.cc b/src/fdt_handler_test.cc
index 0e039bb..ab0b349 100644
--- a/src/fdt_handler_test.cc
+++ b/src/fdt_handler_test.cc
@@ -1,13 +1,18 @@
 extern "C" {
 #include "hf/fdt_handler.h"
 
+#include "hf/alloc.h"
 #include "hf/boot_params.h"
 }
 
+#include <memory>
+
 #include <gmock/gmock.h>
 
 using ::testing::Eq;
 
+static constexpr size_t TEST_HEAP_SIZE = PAGE_SIZE * 10;
+
 /*
  * /dts-v1/;
  *
@@ -35,7 +40,7 @@
  * | xxd -i
  */
 
-static const uint8_t test_dtb[] = {
+static constexpr uint8_t test_dtb[] = {
 	0xd0, 0x0d, 0xfe, 0xed, 0x00, 0x00, 0x01, 0x7f, 0x00, 0x00, 0x00, 0x38,
 	0x00, 0x00, 0x01, 0x30, 0x00, 0x00, 0x00, 0x28, 0x00, 0x00, 0x00, 0x11,
 	0x00, 0x00, 0x00, 0x10, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x4f,
@@ -71,7 +76,11 @@
 
 TEST(fdt, get_boot_params)
 {
-	struct boot_params params;
+	std::unique_ptr<uint8_t[]> test_heap(new uint8_t[TEST_HEAP_SIZE]);
+	halloc_init((size_t)test_heap.get(), TEST_HEAP_SIZE);
+	ASSERT_TRUE(mm_init());
+
+	struct boot_params params = {};
 	EXPECT_TRUE(
 		fdt_get_boot_params(pa_init((uintpaddr_t)&test_dtb), &params));
 	EXPECT_THAT(params.mem_ranges_count, Eq(3));
diff --git a/src/mock/BUILD.gn b/src/mock/BUILD.gn
deleted file mode 100644
index 594db5c..0000000
--- a/src/mock/BUILD.gn
+++ /dev/null
@@ -1,5 +0,0 @@
-source_set("mm") {
-  sources = [
-    "mm.c",
-  ]
-}
diff --git a/src/mock/mm.c b/src/mock/mm.c
deleted file mode 100644
index 3cc0f1f..0000000
--- a/src/mock/mm.c
+++ /dev/null
@@ -1,68 +0,0 @@
-#include "hf/mm.h"
-
-#include "hf/dlog.h"
-
-bool mm_init(void)
-{
-	return true;
-}
-
-bool mm_cpu_init(void)
-{
-	return true;
-}
-
-void mm_defrag(void)
-{
-}
-
-void *mm_identity_map(paddr_t begin, paddr_t end, int mode)
-{
-	(void)end;
-	(void)mode;
-	return ptr_from_va(va_from_pa(begin));
-}
-
-bool mm_unmap(paddr_t begin, paddr_t end, int mode)
-{
-	(void)begin;
-	(void)end;
-	(void)mode;
-	return false;
-}
-
-bool mm_vm_translate(struct mm_ptable *t, ipaddr_t ipa, paddr_t *pa)
-{
-	(void)t;
-	(void)ipa;
-	(void)pa;
-	return false;
-}
-
-bool mm_ptable_init(struct mm_ptable *t, int mode)
-{
-	(void)t;
-	(void)mode;
-	return true;
-}
-
-bool mm_vm_identity_map(struct mm_ptable *t, paddr_t begin, paddr_t end,
-			int mode, ipaddr_t *ipa)
-{
-	(void)t;
-	(void)end;
-	(void)mode;
-	*ipa = ipa_from_pa(begin);
-
-	return true;
-}
-
-bool mm_vm_identity_map_page(struct mm_ptable *t, paddr_t begin, int mode,
-			     ipaddr_t *ipa)
-{
-	(void)t;
-	(void)mode;
-	*ipa = ipa_from_pa(begin);
-
-	return true;
-}