Merge "Add a test Fibonacci extension"

1 files changed, 8 insertions, 1 deletions

diff --git a/nn/runtime/test/TestExtensions.cpp b/nn/runtime/test/TestExtensions.cpp
index b801c08cd..7a228d290 100644
--- a/nn/runtime/test/TestExtensions.cpp
+++ b/nn/runtime/test/TestExtensions.cpp
@@ -21,11 +21,13 @@
 #include "NeuralNetworks.h"
 #include "NeuralNetworksExtensions.h"
 #include "SampleDriver.h"
+#include "TypeManager.h"
 
 namespace {
 
 using DeviceManager = ::android::nn::DeviceManager;
 using SampleDriver = ::android::nn::sample_driver::SampleDriver;
+using TypeManager = ::android::nn::TypeManager;
 
 const char* kTestDriverName = "extensions-test-driver";
 const char* kTestExtension1 = "vendor.test.one";
@@ -67,11 +69,16 @@ class ExtensionsTest : public ::testing::Test {
         }
 
         DeviceManager::get()->forTest_registerDevice(kTestDriverName, new TestDriver());
+        // Discover extensions provided by registered devices.
+        TypeManager::get()->forTest_reset();
         mDevice = getDeviceByName(kTestDriverName);
         ASSERT_NE(mDevice, nullptr);
     }
 
-    virtual void TearDown() { DeviceManager::get()->forTest_reInitializeDeviceList(); }
+    virtual void TearDown() {
+        DeviceManager::get()->forTest_reInitializeDeviceList();
+        TypeManager::get()->forTest_reset();
+    }
 
     ANeuralNetworksDevice* getDeviceByName(const std::string& name) {
         ANeuralNetworksDevice* result = nullptr;