Make DriverDevice::initialize return false for any error

Prior to this CL, DriverDevice::initialize attempted to gracefully
recover from the case when the NN HAL driver fails to return
getNumberOfNeededCacheFiles or when the needed number of cache files
exceeds the specified max number. It recovers gracefully by defaulting
the needed number of cache files to 0, indicating that caching is not
supported. This CL changes the behavior such that this is considered
an initialization failure, causing the calling code not to use that
DriverDevice object.

This CL also changes the TestCompilationCaching to match this new
initialize behavior.

Fixes: 132322236
Test: mma
Test: NeuralNetworksTest_static
Test: CtsNNAPITestCases
Change-Id: I97e40be2553e06f4c97727184286000af2529d45

1 files changed, 127 insertions, 62 deletions

diff --git a/nn/runtime/test/TestCompilationCaching.cpp b/nn/runtime/test/TestCompilationCaching.cpp
index c061e853d..bce2836d2 100644
--- a/nn/runtime/test/TestCompilationCaching.cpp
+++ b/nn/runtime/test/TestCompilationCaching.cpp
@@ -14,12 +14,14 @@
  * limitations under the License.
  */
 
+#include <android-base/scopeguard.h>
 #include <gtest/gtest.h>
 
 #include <cstdlib>
 #include <filesystem>
 #include <numeric>
 #include <string>
+#include <string_view>
 #include <tuple>
 #include <vector>
 
@@ -48,13 +50,32 @@ namespace {
 
 enum class HasCalledPrepareModel { NO, WITHOUT_CACHING, WITH_CACHING };
 
+// Print HasCalledPrepareModel enum for better GTEST failure messages
+std::ostream& operator<<(std::ostream& os, HasCalledPrepareModel hasCalledPrepareModel) {
+    switch (hasCalledPrepareModel) {
+        case HasCalledPrepareModel::NO:
+            return os << "NO";
+        case HasCalledPrepareModel::WITHOUT_CACHING:
+            return os << "WITHOUT_CACHING";
+        case HasCalledPrepareModel::WITH_CACHING:
+            return os << "WITH_CACHING";
+    }
+    CHECK(false) << "HasCalledPrepareModel print called with invalid code "
+                 << static_cast<int>(hasCalledPrepareModel);
+    return os;
+}
+
+// Whether the driver is expected to be registered because it can pass initialization.
+bool canDeviceBeRegistered(ErrorStatus error, uint32_t numModelCache, uint32_t numDataCache) {
+    constexpr uint32_t maxNumCacheFiles =
+            static_cast<uint32_t>(Constant::MAX_NUMBER_OF_CACHE_FILES);
+    return error == ErrorStatus::NONE && numModelCache <= maxNumCacheFiles &&
+           numDataCache <= maxNumCacheFiles;
+}
+
 // Whether the driver supports caching based on the returns from getNumberOfCacheFilesNeeded.
-bool isCachingSupportedAndNoError(ErrorStatus error, uint32_t numModelCache,
-                                  uint32_t numDataCache) {
-    return error == ErrorStatus::NONE &&
-           numModelCache <= static_cast<uint32_t>(Constant::MAX_NUMBER_OF_CACHE_FILES) &&
-           numDataCache <= static_cast<uint32_t>(Constant::MAX_NUMBER_OF_CACHE_FILES) &&
-           (numModelCache != 0 || numDataCache != 0);
+bool isCachingSupported(uint32_t numModelCache, uint32_t numDataCache) {
+    return numModelCache != 0 || numDataCache != 0;
 }
 
 // This is an IDevice for testing purposes which overrides several methods from sample driver:
@@ -99,9 +120,10 @@ class CachingDriver : public sample_driver::SampleDriver {
     };
 
    public:
-    CachingDriver(const char* name, ErrorStatus errorStatusGetNumCacheFiles, uint32_t numModelCache,
-                  uint32_t numDataCache, ErrorStatus errorStatusPrepareFromCache)
-        : SampleDriver(name),
+    CachingDriver(std::string_view name, ErrorStatus errorStatusGetNumCacheFiles,
+                  uint32_t numModelCache, uint32_t numDataCache,
+                  ErrorStatus errorStatusPrepareFromCache)
+        : SampleDriver(name.data()),
           mErrorStatusGetNumCacheFiles(errorStatusGetNumCacheFiles),
           mNumModelCache(numModelCache),
           mNumDataCache(numDataCache),
@@ -181,8 +203,7 @@ class CachingDriver : public sample_driver::SampleDriver {
    private:
     // Checks the number of cache files passed to the driver from runtime.
     void checkNumberOfCacheHandles(size_t modelCache, size_t dataCache) {
-        if (isCachingSupportedAndNoError(mErrorStatusGetNumCacheFiles, mNumModelCache,
-                                         mNumDataCache)) {
+        if (isCachingSupported(mNumModelCache, mNumDataCache)) {
             if (modelCache != 0 || dataCache != 0) {
                 ASSERT_EQ(modelCache, mNumModelCache);
                 ASSERT_EQ(dataCache, mNumDataCache);
@@ -239,12 +260,69 @@ void CreateBroadcastAddModel(test_wrapper::Model* model) {
     ASSERT_EQ(model->finish(), Result::NO_ERROR);
 }
 
-// Test model compilation with a driver parameterized with
+void getDeviceWithName(std::string_view deviceName, const ANeuralNetworksDevice** outputDevice) {
+    uint32_t numDevices = 0;
+    ASSERT_EQ(ANeuralNetworks_getDeviceCount(&numDevices), ANEURALNETWORKS_NO_ERROR);
+    EXPECT_GE(numDevices, (uint32_t)1);
+
+    int numMatchingDevices = 0;
+    for (uint32_t i = 0; i < numDevices; i++) {
+        ANeuralNetworksDevice* device = nullptr;
+        ASSERT_EQ(ANeuralNetworks_getDevice(i, &device), ANEURALNETWORKS_NO_ERROR);
+
+        const char* buffer = nullptr;
+        ASSERT_EQ(ANeuralNetworksDevice_getName(device, &buffer), ANEURALNETWORKS_NO_ERROR);
+        if (deviceName == buffer) {
+            *outputDevice = device;
+            numMatchingDevices++;
+        }
+    }
+
+    EXPECT_LE(numMatchingDevices, 1);
+}
+
+// Test device registration with a driver parameterized with
 // - ErrorStatus returning from getNumberOfCacheFilesNeeded
 // - Number of model cache files returning from getNumberOfCacheFilesNeeded
 // - Number of data cache files returning from getNumberOfCacheFilesNeeded
+using DeviceRegistrationTestParam = std::tuple<ErrorStatus, uint32_t, uint32_t>;
+
+class DeviceRegistrationTest : public ::testing::TestWithParam<DeviceRegistrationTestParam> {
+   protected:
+    static constexpr std::string_view kDeviceName = "deviceTestCompilationCaching";
+    const ErrorStatus kErrorStatusGetNumCacheFiles = std::get<0>(GetParam());
+    const uint32_t kNumModelCache = std::get<1>(GetParam());
+    const uint32_t kNumDataCache = std::get<2>(GetParam());
+    const sp<CachingDriver> kDriver =
+            new CachingDriver(kDeviceName, kErrorStatusGetNumCacheFiles, kNumModelCache,
+                              kNumDataCache, ErrorStatus::NONE);
+};
+
+TEST_P(DeviceRegistrationTest, CachingFailure) {
+    if (DeviceManager::get()->getUseCpuOnly()) {
+        return;
+    }
+
+    DeviceManager::get()->forTest_registerDevice(kDeviceName.data(), kDriver);
+    const auto cleanup = android::base::make_scope_guard(
+            [] { DeviceManager::get()->forTest_reInitializeDeviceList(); });
+
+    // get device
+    const ANeuralNetworksDevice* device = nullptr;
+    getDeviceWithName(kDeviceName, &device);
+
+    // check if device registeration matches expectations
+    const bool isDeviceRegistered = (device != nullptr);
+    const bool expectDeviceToBeRegistered =
+            canDeviceBeRegistered(kErrorStatusGetNumCacheFiles, kNumModelCache, kNumDataCache);
+    ASSERT_EQ(isDeviceRegistered, expectDeviceToBeRegistered);
+}
+
+// Test model compilation with a driver parameterized with
+// - Number of model cache files returning from getNumberOfCacheFilesNeeded
+// - Number of data cache files returning from getNumberOfCacheFilesNeeded
 // - ErrorStatus returning from prepareModelFromCache
-using CompilationCachingTestParam = std::tuple<ErrorStatus, uint32_t, uint32_t, ErrorStatus>;
+using CompilationCachingTestParam = std::tuple<uint32_t, uint32_t, ErrorStatus>;
 
 class CompilationCachingTest : public ::testing::TestWithParam<CompilationCachingTestParam> {
    protected:
@@ -254,9 +332,6 @@ class CompilationCachingTest : public ::testing::TestWithParam<CompilationCachin
         ASSERT_NE(cacheDir, nullptr);
         mCacheDir = cacheDir;
         CreateBroadcastAddModel(&mModel);
-        mToken = std::vector<uint8_t>(ANEURALNETWORKS_BYTE_SIZE_OF_CACHE_TOKEN, 0);
-        mIsCachingSupportedAndNoError = isCachingSupportedAndNoError(kErrorStatusGetNumCacheFiles,
-                                                                     kNumModelCache, kNumDataCache);
     }
 
     virtual void TearDown() override {
@@ -266,38 +341,26 @@ class CompilationCachingTest : public ::testing::TestWithParam<CompilationCachin
     }
 
     void compileModel(const sp<CachingDriver>& driver, bool withToken) {
-        DeviceManager::get()->forTest_registerDevice(kDeviceName, driver);
-
-        // Make device list including only a single driver device.
-        uint32_t numDevices = 0;
-        EXPECT_EQ(ANeuralNetworks_getDeviceCount(&numDevices), ANEURALNETWORKS_NO_ERROR);
-        EXPECT_GE(numDevices, (uint32_t)1);
-        std::vector<ANeuralNetworksDevice*> devices;
-        for (uint32_t i = 0; i < numDevices; i++) {
-            ANeuralNetworksDevice* device = nullptr;
-            EXPECT_EQ(ANeuralNetworks_getDevice(i, &device), ANEURALNETWORKS_NO_ERROR);
-            const char* buffer = nullptr;
-            int result = ANeuralNetworksDevice_getName(device, &buffer);
-            if (result == ANEURALNETWORKS_NO_ERROR && strcmp(buffer, kDeviceName) == 0) {
-                devices.push_back(device);
-                break;
-            }
-        }
-        ASSERT_EQ(devices.size(), 1u);
+        DeviceManager::get()->forTest_registerDevice(kDeviceName.data(), driver);
+        const auto cleanup = android::base::make_scope_guard(
+                [] { DeviceManager::get()->forTest_reInitializeDeviceList(); });
+
+        // Get a handle to the single driver device matching kDeviceName.
+        const ANeuralNetworksDevice* device = nullptr;
+        getDeviceWithName(kDeviceName, &device);
+        ASSERT_NE(device, nullptr);
 
         // Compile the model with the device.
         ANeuralNetworksCompilation* compilation = nullptr;
-        ASSERT_EQ(ANeuralNetworksCompilation_createForDevices(mModel.getHandle(), devices.data(),
-                                                              devices.size(), &compilation),
+        ASSERT_EQ(ANeuralNetworksCompilation_createForDevices(mModel.getHandle(), &device, 1,
+                                                              &compilation),
                   ANEURALNETWORKS_NO_ERROR);
         if (withToken) {
             ASSERT_EQ(ANeuralNetworksCompilation_setCaching(compilation, mCacheDir.c_str(),
-                                                            mToken.data()),
+                                                            kToken.data()),
                       ANEURALNETWORKS_NO_ERROR);
         }
         ASSERT_EQ(ANeuralNetworksCompilation_finish(compilation), ANEURALNETWORKS_NO_ERROR);
-
-        DeviceManager::get()->forTest_reInitializeDeviceList();
     }
 
     void createCache() {
@@ -306,31 +369,29 @@ class CompilationCachingTest : public ::testing::TestWithParam<CompilationCachin
         compileModel(driver, /*withToken=*/true);
     }
 
-    static constexpr char kDeviceName[] = "deviceTestCompilationCaching";
-    const ErrorStatus kErrorStatusGetNumCacheFiles = std::get<0>(GetParam());
-    const uint32_t kNumModelCache = std::get<1>(GetParam());
-    const uint32_t kNumDataCache = std::get<2>(GetParam());
-    const ErrorStatus kErrorStatusPrepareFromCache = std::get<3>(GetParam());
-    bool mIsCachingSupportedAndNoError;
+    static constexpr std::string_view kDeviceName = "deviceTestCompilationCaching";
+    const uint32_t kNumModelCache = std::get<0>(GetParam());
+    const uint32_t kNumDataCache = std::get<1>(GetParam());
+    const ErrorStatus kErrorStatusPrepareFromCache = std::get<2>(GetParam());
+    const bool kIsCachingSupported = isCachingSupported(kNumModelCache, kNumDataCache);
     test_wrapper::Model mModel;
     std::string mCacheDir;
-    std::vector<uint8_t> mToken;
+    const CacheToken kToken{};
 };
 
 TEST_P(CompilationCachingTest, TokenProvidedAndCacheNotExist) {
     if (DeviceManager::get()->getUseCpuOnly()) {
         return;
     }
-    sp<CachingDriver> driver =
-            new CachingDriver(kDeviceName, kErrorStatusGetNumCacheFiles, kNumModelCache,
-                              kNumDataCache, kErrorStatusPrepareFromCache);
+    sp<CachingDriver> driver = new CachingDriver(kDeviceName, ErrorStatus::NONE, kNumModelCache,
+                                                 kNumDataCache, kErrorStatusPrepareFromCache);
     compileModel(driver, /*withToken=*/true);
 
     // When cache file does not exist, the runtime should never call prepareModelFromCache.
-    EXPECT_EQ(driver->hasCalledPrepareModelFromCache(), false);
+    EXPECT_FALSE(driver->hasCalledPrepareModelFromCache());
 
     // The runtime should call prepareModel_1_2. It should request caching iff caching supported.
-    EXPECT_EQ(driver->hasCalledPrepareModel(), mIsCachingSupportedAndNoError
+    EXPECT_EQ(driver->hasCalledPrepareModel(), kIsCachingSupported
                                                        ? HasCalledPrepareModel::WITH_CACHING
                                                        : HasCalledPrepareModel::WITHOUT_CACHING);
 }
@@ -340,16 +401,15 @@ TEST_P(CompilationCachingTest, TokenProvidedAndCacheExist) {
         return;
     }
     createCache();
-    sp<CachingDriver> driver =
-            new CachingDriver(kDeviceName, kErrorStatusGetNumCacheFiles, kNumModelCache,
-                              kNumDataCache, kErrorStatusPrepareFromCache);
+    sp<CachingDriver> driver = new CachingDriver(kDeviceName, ErrorStatus::NONE, kNumModelCache,
+                                                 kNumDataCache, kErrorStatusPrepareFromCache);
     compileModel(driver, /*withToken=*/true);
 
     // When cache files exist, the runtime should call prepareModelFromCache iff caching supported.
-    EXPECT_EQ(driver->hasCalledPrepareModelFromCache(), mIsCachingSupportedAndNoError);
+    EXPECT_EQ(driver->hasCalledPrepareModelFromCache(), kIsCachingSupported);
 
     HasCalledPrepareModel expectHasCalledPrepareModel;
-    if (mIsCachingSupportedAndNoError) {
+    if (kIsCachingSupported) {
         if (kErrorStatusPrepareFromCache == ErrorStatus::NONE) {
             // The runtime should not call prepareModel_1_2 iff caching supported and
             // prepareModelFromCache succeeds.
@@ -370,14 +430,13 @@ TEST_P(CompilationCachingTest, TokenNotProvided) {
     if (DeviceManager::get()->getUseCpuOnly()) {
         return;
     }
-    sp<CachingDriver> driver =
-            new CachingDriver(kDeviceName, kErrorStatusGetNumCacheFiles, kNumModelCache,
-                              kNumDataCache, kErrorStatusPrepareFromCache);
+    sp<CachingDriver> driver = new CachingDriver(kDeviceName, ErrorStatus::NONE, kNumModelCache,
+                                                 kNumDataCache, kErrorStatusPrepareFromCache);
     compileModel(driver, /*withToken=*/false);
 
     // When no NDK token is provided by the client, the runtime should never call
     // prepareModelFromCache or request caching with prepareModel_1_2.
-    EXPECT_EQ(driver->hasCalledPrepareModelFromCache(), false);
+    EXPECT_FALSE(driver->hasCalledPrepareModelFromCache());
     EXPECT_EQ(driver->hasCalledPrepareModel(), HasCalledPrepareModel::WITHOUT_CACHING);
 }
 
@@ -386,12 +445,18 @@ static const auto kErrorStatusGetNumCacheFilesChoices =
 static const auto kNumCacheChoices =
         testing::Values(0ul, 1ul, static_cast<uint32_t>(Constant::MAX_NUMBER_OF_CACHE_FILES),
                         static_cast<uint32_t>(Constant::MAX_NUMBER_OF_CACHE_FILES) + 1);
+static const auto kNumValidCacheChoices =
+        testing::Values(0ul, 1ul, static_cast<uint32_t>(Constant::MAX_NUMBER_OF_CACHE_FILES));
 static const auto kErrorStatusPrepareFromCacheChoices =
         testing::Values(ErrorStatus::NONE, ErrorStatus::GENERAL_FAILURE,
                         ErrorStatus::DEVICE_UNAVAILABLE, ErrorStatus::INVALID_ARGUMENT);
 
-INSTANTIATE_TEST_CASE_P(TestCompilationCaching, CompilationCachingTest,
+INSTANTIATE_TEST_CASE_P(TestCompilationCaching, DeviceRegistrationTest,
                         testing::Combine(kErrorStatusGetNumCacheFilesChoices, kNumCacheChoices,
-                                         kNumCacheChoices, kErrorStatusPrepareFromCacheChoices));
+                                         kNumCacheChoices));
+
+INSTANTIATE_TEST_CASE_P(TestCompilationCaching, CompilationCachingTest,
+                        testing::Combine(kNumValidCacheChoices, kNumValidCacheChoices,
+                                         kErrorStatusPrepareFromCacheChoices));
 
 }  // namespace