diff --git a/include/villas/memory_type.h b/include/villas/memory_type.h
index b1167992b..50ca0f5dc 100644
--- a/include/villas/memory_type.h
+++ b/include/villas/memory_type.h
@@ -59,6 +59,3 @@ struct memory_type * memory_ib(struct node *n, struct memory_type *parent);
 struct memory_type * memory_managed(void *ptr, size_t len);
 
 int memory_hugepage_init(int hugepages);
-
-struct memory_type * memory_type_lookup(enum MemoryFlags flags);
-
diff --git a/lib/memory.cpp b/lib/memory.cpp
index fe5f0ccd2..846f892cc 100644
--- a/lib/memory.cpp
+++ b/lib/memory.cpp
@@ -44,9 +44,9 @@ int memory_init(int hugepages)
 
 	info("Initialize memory sub-system: #hugepages=%d", hugepages);
 
-	ret = memory_hugepage_init(hugepages);
-	if (ret)
-		return ret;
+		ret = memory_hugepage_init(hugepages);
+		if (ret)
+			return ret;
 
 	size_t lock = kernel_get_hugepage_size() * hugepages;
 
@@ -142,10 +142,10 @@ int memory_free(void *ptr)
 
 	/* Remove allocation entry */
 	auto iter = allocations.find(ptr);
-  	if (iter == allocations.end())
+	if (iter == allocations.end())
 		return -1;
 
-    allocations.erase(iter);
+	allocations.erase(iter);
 	free(ma);
 
 	return 0;
@@ -156,12 +156,4 @@ struct memory_allocation * memory_get_allocation(void *ptr)
 	return allocations[ptr];
 }
 
-struct memory_type * memory_type_lookup(enum MemoryFlags flags)
-{
-	if ((int) flags & (int) MemoryFlags::HUGEPAGE)
-		return &memory_hugepage;
-	else if ((int) flags & (int) MemoryFlags::HEAP)
-		return &memory_heap;
-	else
-		return nullptr;
-}
+struct memory_type *memory_default = nullptr;
diff --git a/tests/unit/memory.cpp b/tests/unit/memory.cpp
index 6b09cd719..f71c5d68f 100644
--- a/tests/unit/memory.cpp
+++ b/tests/unit/memory.cpp
@@ -37,15 +37,13 @@ extern void init_memory();
 TheoryDataPoints(memory, aligned) = {
 	DataPoints(size_t, 1, 32, 55, 1 << 10, PAGESIZE, HUGEPAGESIZE),
 	DataPoints(size_t, 1, 8, PAGESIZE, PAGESIZE),
-	DataPoints(enum MemoryFlags, MemoryFlags::HEAP, MemoryFlags::HUGEPAGE, MemoryFlags::HUGEPAGE)
+	DataPoints(struct memory_type *, &memory_heap, &memory_mmap_hugetlb, &memory_mmap_hugetlb)
 };
 
-Theory((size_t len, size_t align, enum MemoryFlags memory_type), memory, aligned, .init = init_memory) {
+Theory((size_t len, size_t align, struct memory_type *mt), memory, aligned, .init = init_memory) {
 	int ret;
 	void *ptr;
 
-	struct memory_type *mt = memory_type_lookup(memory_type);
-
 	ptr = memory_alloc_aligned(len, align, mt);
 	cr_assert_not_null(ptr, "Failed to allocate memory");
 
diff --git a/tests/unit/pool.cpp b/tests/unit/pool.cpp
index 1edf50933..1b3d5b538 100644
--- a/tests/unit/pool.cpp
+++ b/tests/unit/pool.cpp
@@ -34,16 +34,16 @@ struct param {
 	int thread_count;
 	int pool_size;
 	size_t block_size;
-	enum MemoryFlags memory_type;
+	struct memory_type *mt;
 };
 
 ParameterizedTestParameters(pool, basic)
 {
 	static struct param params[] = {
-		{ 1,	4096,	150,	MemoryFlags::HEAP },
-		{ 1,	128,	8,	MemoryFlags::HUGEPAGE },
-		{ 1,	4,	8192,	MemoryFlags::HUGEPAGE },
-		{ 1,	1 << 13, 4,	MemoryFlags::HUGEPAGE }
+		{ 1,	4096,	150,	&memory_heap },
+		{ 1,	128,	8,	&memory_mmap_hugetlb },
+		{ 1,	4,	8192,	&memory_mmap_hugetlb },
+		{ 1,	1 << 13, 4,	&memory_mmap_hugetlb }
 	};
 
 	return cr_make_param_array(struct param, params, ARRAY_LEN(params));
@@ -56,9 +56,7 @@ ParameterizedTest(struct param *p, pool, basic, .init = init_memory)
 
 	void *ptr, *ptrs[p->pool_size];
 
-	struct memory_type *mt = memory_type_lookup(p->memory_type);
-
-	ret = pool_init(&pool, p->pool_size, p->block_size, mt);
+	ret = pool_init(&pool, p->pool_size, p->block_size, p->mt);
 	cr_assert_eq(ret, 0, "Failed to create pool");
 
 	ptr = pool_get(&pool);
diff --git a/tests/unit/queue.cpp b/tests/unit/queue.cpp
index 714e8b357..652d42431 100644
--- a/tests/unit/queue.cpp
+++ b/tests/unit/queue.cpp
@@ -54,7 +54,7 @@ struct param {
 	int thread_count;
 	bool many;
 	int batch_size;
-	enum MemoryFlags memory_type;
+	struct memory_type *mt;
 	volatile int start;
 	struct queue queue;
 };
@@ -274,35 +274,35 @@ ParameterizedTestParameters(queue, multi_threaded)
 			.thread_count = 32,
 			.many = true,
 			.batch_size = 10,
-			.memory_type = MemoryFlags::HEAP
+			.mt = &memory_heap
 		}, {
 			.iter_count = 1 << 8,
 			.queue_size = 1 << 9,
 			.thread_count = 4,
 			.many = true,
 			.batch_size = 100,
-			.memory_type = MemoryFlags::HEAP
+			.mt = &memory_heap
 		}, {
 			.iter_count = 1 << 16,
 			.queue_size = 1 << 14,
 			.thread_count = 16,
 			.many = true,
 			.batch_size = 100,
-			.memory_type = MemoryFlags::HEAP
+			.mt = &memory_heap
 		}, {
 			.iter_count = 1 << 8,
 			.queue_size = 1 << 9,
 			.thread_count = 4,
 			.many = true,
 			.batch_size = 10,
-			.memory_type = MemoryFlags::HEAP
+			.mt = &memory_heap
 		}, {
 			.iter_count = 1 << 16,
 			.queue_size = 1 << 9,
 			.thread_count = 16,
 			.many = false,
 			.batch_size = 10,
-			.memory_type = MemoryFlags::HUGEPAGE
+			.mt = &memory_mmap_hugetlb
 		}
 	};
 
@@ -321,9 +321,7 @@ ParameterizedTest(struct param *p, queue, multi_threaded, .timeout = 20, .init =
 	p->start = 0;
 	p->queue.state = ATOMIC_VAR_INIT(State::DESTROYED);
 
-	struct memory_type *mt = memory_type_lookup(p->memory_type);
-
-	ret = queue_init(&p->queue, p->queue_size, mt);
+	ret = queue_init(&p->queue, p->queue_size, p->mt);
 	cr_assert_eq(ret, 0, "Failed to create queue");
 
 	uint64_t start_tsc_time, end_tsc_time;