diff --git a/oak_functions_service/benches/wasm_benchmark.rs b/oak_functions_service/benches/wasm_benchmark.rs
index 00e156bc827..78e958b7083 100644
--- a/oak_functions_service/benches/wasm_benchmark.rs
+++ b/oak_functions_service/benches/wasm_benchmark.rs
@@ -65,7 +65,8 @@ fn bench_invoke_lookup(c: &mut Criterion) {
     let test_data = create_test_data(0, MAX_DATA_SIZE);
     test_state
         .lookup_data_manager
-        .extend_next_lookup_data(test_data.clone());
+        .extend_next_lookup_data(test_data.clone())
+        .unwrap();
     test_state.lookup_data_manager.finish_next_lookup_data();
 
     c.bench_function("lookup wasm", |b| {
@@ -121,14 +122,16 @@ fn bench_invoke_lookup_multi(c: &mut Criterion) {
     let test_data = create_test_data(0, MAX_DATA_SIZE);
     test_state_wasmi
         .lookup_data_manager
-        .extend_next_lookup_data(test_data.clone());
+        .extend_next_lookup_data(test_data.clone())
+        .unwrap();
     test_state_wasmi
         .lookup_data_manager
         .finish_next_lookup_data();
 
     test_state_wasmtime
         .lookup_data_manager
-        .extend_next_lookup_data(test_data.clone());
+        .extend_next_lookup_data(test_data.clone())
+        .unwrap();
     test_state_wasmtime
         .lookup_data_manager
         .finish_next_lookup_data();
diff --git a/oak_functions_service/src/lookup.rs b/oak_functions_service/src/lookup.rs
index 197e6ef86f1..93ee5e80fbb 100644
--- a/oak_functions_service/src/lookup.rs
+++ b/oak_functions_service/src/lookup.rs
@@ -18,17 +18,17 @@ use alloc::{
     format,
     string::{String, ToString},
     sync::Arc,
-    vec::Vec,
 };
 
 use bytes::Bytes;
+use hashbrown::HashMap;
 use log::{info, Level};
 use spinning_top::{RwSpinlock, Spinlock};
 
 use crate::{logger::OakLogger, lookup_htbl::LookupHtbl};
 
 // This is used to pass data to the lookup data module for tests.
-pub type Data = Vec<(Bytes, Bytes)>;
+pub type Data = HashMap<Bytes, Bytes>;
 
 // Data maintains the invariant on lookup data to have [at most one
 // value](https://github.com/project-oak/oak/tree/main/oak/oak_functions_service/README.md#invariant-at-most-one-value)
@@ -107,7 +107,7 @@ impl LookupDataManager {
     }
 
     /// Creates an instance of LookupData populated with the given entries.
-    pub fn for_test(data: Vec<(Bytes, Bytes)>, logger: Arc<dyn OakLogger>) -> Self {
+    pub fn for_test(data: Data, logger: Arc<dyn OakLogger>) -> Self {
         let test_manager = Self::new_empty(logger);
         test_manager.reserve(data.len() as u64).unwrap();
         test_manager.extend_next_lookup_data(data).unwrap();
@@ -233,7 +233,7 @@ pub fn format_bytes(v: &[u8]) -> String {
 
 #[cfg(test)]
 mod tests {
-    use alloc::vec;
+    use alloc::{vec, vec::Vec};
 
     use super::*;