Merge "Allow getPsdsExtension to return nullptr" into stage-aosp-sc-ts-dev
diff --git a/neuralnetworks/aidl/vts/functional/Utils.cpp b/neuralnetworks/aidl/vts/functional/Utils.cpp
index 9af362e..325a436 100644
--- a/neuralnetworks/aidl/vts/functional/Utils.cpp
+++ b/neuralnetworks/aidl/vts/functional/Utils.cpp
@@ -153,26 +153,19 @@
             .stride = size,
     };
 
-    ASSERT_EQ(AHardwareBuffer_allocate(&desc, &mAhwb), 0);
-    ASSERT_NE(mAhwb, nullptr);
+    AHardwareBuffer* ahwb = nullptr;
+    ASSERT_EQ(AHardwareBuffer_allocate(&desc, &ahwb), 0);
+    ASSERT_NE(ahwb, nullptr);
 
-    const auto sharedMemory =
-            nn::createSharedMemoryFromAHWB(mAhwb, /*takeOwnership=*/false).value();
-    mMapping = nn::map(sharedMemory).value();
+    mMemory = nn::createSharedMemoryFromAHWB(ahwb, /*takeOwnership=*/true).value();
+    mMapping = nn::map(mMemory).value();
     mPtr = static_cast<uint8_t*>(std::get<void*>(mMapping.pointer));
     CHECK_NE(mPtr, nullptr);
-    mAidlMemory = utils::convert(sharedMemory).value();
+    mAidlMemory = utils::convert(mMemory).value();
 
     mIsValid = true;
 }
 
-TestBlobAHWB::~TestBlobAHWB() {
-    if (mAhwb) {
-        AHardwareBuffer_unlock(mAhwb, nullptr);
-        AHardwareBuffer_release(mAhwb);
-    }
-}
-
 std::string gtestCompliantName(std::string name) {
     // gtest test names must only contain alphanumeric characters
     std::replace_if(
diff --git a/neuralnetworks/aidl/vts/functional/Utils.h b/neuralnetworks/aidl/vts/functional/Utils.h
index 9dd7359..ca81418 100644
--- a/neuralnetworks/aidl/vts/functional/Utils.h
+++ b/neuralnetworks/aidl/vts/functional/Utils.h
@@ -102,11 +102,10 @@
     // The constructor calls initialize, which constructs the memory resources. This is a
     // workaround that gtest macros cannot be used directly in a constructor.
     TestBlobAHWB(uint32_t size) { initialize(size); }
-    ~TestBlobAHWB();
 
   private:
     void initialize(uint32_t size);
-    AHardwareBuffer* mAhwb = nullptr;
+    nn::SharedMemory mMemory;
     nn::Mapping mMapping;
 };