diff --git a/tests/memory/1.0/Android.bp b/tests/memory/1.0/Android.bp
index cd97fe4..5038664 100644
--- a/tests/memory/1.0/Android.bp
+++ b/tests/memory/1.0/Android.bp
@@ -7,6 +7,8 @@
         "IMemoryTest.hal",
     ],
     interfaces: [
+        "android.hidl.memory.token@1.0",
+        "android.hidl.memory.block@1.0",
         "android.hidl.base@1.0",
     ],
     gen_java: false,
diff --git a/tests/memory/1.0/IMemoryTest.hal b/tests/memory/1.0/IMemoryTest.hal
index 4d6de3f..b20859c 100644
--- a/tests/memory/1.0/IMemoryTest.hal
+++ b/tests/memory/1.0/IMemoryTest.hal
@@ -16,7 +16,13 @@
 
 package android.hardware.tests.memory@1.0;
 
+import android.hidl.memory.token@1.0::IMemoryToken;
+import android.hidl.memory.block@1.0::MemoryBlock;
+
 interface IMemoryTest {
     haveSomeMemory(memory mem) generates(memory mem);
     fillMemory(memory memory_in, uint8_t filler);
+    haveSomeMemoryBlock(MemoryBlock blk) generates(MemoryBlock blk);
+    set(memory mem) generates();
+    get()generates(IMemoryToken token);
 };
diff --git a/tests/memory/1.0/default/Android.bp b/tests/memory/1.0/default/Android.bp
index 985183b..3f13634 100644
--- a/tests/memory/1.0/default/Android.bp
+++ b/tests/memory/1.0/default/Android.bp
@@ -28,6 +28,7 @@
         "liblog",
         "libutils",
         "android.hidl.memory@1.0",
+        "android.hidl.memory.token@1.0",
     ],
 
     // These are static libs only for testing purposes and portability. Shared
diff --git a/tests/memory/1.0/default/MemoryTest.cpp b/tests/memory/1.0/default/MemoryTest.cpp
index 37a2a60..5cbf67e 100644
--- a/tests/memory/1.0/default/MemoryTest.cpp
+++ b/tests/memory/1.0/default/MemoryTest.cpp
@@ -20,6 +20,7 @@
 
 #include <log/log.h>
 
+#include <hidlmemory/HidlMemoryToken.h>
 #include <hidlmemory/mapping.h>
 
 #include <android/hidl/memory/1.0/IMemory.h>
@@ -60,6 +61,22 @@
     return Void();
 }
 
+Return<void> Memory::haveSomeMemoryBlock(const MemoryBlock& blk, haveSomeMemoryBlock_cb _hidl_cb) {
+    _hidl_cb(blk);
+    return Void();
+}
+
+Return<void> Memory::set(const hidl_memory& mem) {
+    sp<HidlMemory> hidlMem = HidlMemory::getInstance(mem);
+    if (hidlMem->valid()) {
+        mSavedMemoryToken = new HidlMemoryToken(hidlMem);
+    }
+    return Void();
+}
+
+Return<sp<IMemoryToken>> Memory::get() {
+    return mSavedMemoryToken;
+}
 
 IMemoryTest* HIDL_FETCH_IMemoryTest(const char* /* name */) {
     return new Memory();
diff --git a/tests/memory/1.0/default/MemoryTest.h b/tests/memory/1.0/default/MemoryTest.h
index 0d903f1..4c2ed36 100644
--- a/tests/memory/1.0/default/MemoryTest.h
+++ b/tests/memory/1.0/default/MemoryTest.h
@@ -20,6 +20,7 @@
 #include <android/hardware/tests/memory/1.0/IMemoryTest.h>
 #include <hidl/MQDescriptor.h>
 #include <hidl/Status.h>
+#include <hidlmemory/HidlMemoryToken.h>
 
 namespace android {
 namespace hardware {
@@ -28,14 +29,17 @@
 namespace V1_0 {
 namespace implementation {
 
-using ::android::hardware::tests::memory::V1_0::IMemoryTest;
+using ::android::sp;
 using ::android::hardware::hidl_array;
 using ::android::hardware::hidl_memory;
 using ::android::hardware::hidl_string;
 using ::android::hardware::hidl_vec;
+using ::android::hardware::HidlMemoryToken;
 using ::android::hardware::Return;
 using ::android::hardware::Void;
-using ::android::sp;
+using ::android::hardware::tests::memory::V1_0::IMemoryTest;
+using ::android::hidl::memory::block::V1_0::MemoryBlock;
+using ::android::hidl::memory::token::V1_0::IMemoryToken;
 
 struct Memory : public IMemoryTest {
     // Methods from ::android::hardware::tests::memory::V1_0::IMemoryTest follow.
@@ -43,6 +47,14 @@
 
     Return<void> fillMemory(const hidl_memory& memory_in, uint8_t filler) override;
 
+    Return<void> haveSomeMemoryBlock(const MemoryBlock& blk,
+                                     haveSomeMemoryBlock_cb _hidl_cb) override;
+    Return<void> set(const ::android::hardware::hidl_memory& mem) override;
+
+    Return<sp<IMemoryToken>> get() override;
+
+   protected:
+    sp<HidlMemoryToken> mSavedMemoryToken;
 };
 
 extern "C" IMemoryTest* HIDL_FETCH_IMemoryTest(const char* name);