1 files changed, 83 insertions, 1 deletions
diff --git a/neuralnetworks/1.2/utils/test/PreparedModelTest.cpp b/neuralnetworks/1.2/utils/test/PreparedModelTest.cpp
index 5062ac9a97..d297b1a417 100644
--- a/neuralnetworks/1.2/utils/test/PreparedModelTest.cpp
+++ b/neuralnetworks/1.2/utils/test/PreparedModelTest.cpp
@@ -16,6 +16,8 @@
 
 #include "MockPreparedModel.h"
 
+#include "MockBurstContext.h"
+
 #include <android/hardware/neuralnetworks/1.2/IExecutionCallback.h>
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
@@ -67,6 +69,17 @@ auto makeExecuteAsynchronously(V1_0::ErrorStatus launchStatus, V1_0::ErrorStatus
         return launchStatus;
     };
 }
+auto makeConfigureExecutionBurstReturn(V1_0::ErrorStatus status,
+                                       const sp<MockBurstContext>& burstContext) {
+    return [status, burstContext](
+                   const sp<V1_2::IBurstCallback>& /*callback*/,
+                   const MQDescriptorSync<V1_2::FmqRequestDatum>& /*requestChannel*/,
+                   const MQDescriptorSync<V1_2::FmqResultDatum>& /*resultChannel*/,
+                   V1_2::IPreparedModel::configureExecutionBurst_cb cb) -> hardware::Return<void> {
+        cb(status, burstContext);
+        return hardware::Void();
+    };
+}
 
 std::function<hardware::Status()> makeTransportFailure(status_t status) {
     return [status] { return hardware::Status::fromStatusT(status); };
@@ -321,7 +334,76 @@ TEST(PreparedModelTest, executeFencedNotSupported) {
     EXPECT_EQ(result.error().code, nn::ErrorStatus::GENERAL_FAILURE);
 }
 
-// TODO: test burst execution if/when it is added to nn::IPreparedModel.
+TEST(PreparedModelTest, configureExecutionBurst) {
+    // setup test
+    const auto mockPreparedModel = MockPreparedModel::create();
+    const auto mockBurstContext = sp<MockBurstContext>::make();
+    EXPECT_CALL(*mockPreparedModel, configureExecutionBurst(_, _, _, _))
+            .Times(1)
+            .WillOnce(makeConfigureExecutionBurstReturn(V1_0::ErrorStatus::NONE, mockBurstContext));
+    const auto preparedModel =
+            PreparedModel::create(mockPreparedModel, /*executeSynchronously=*/true).value();
+
+    // run test
+    const auto result = preparedModel->configureExecutionBurst();
+
+    // verify result
+    ASSERT_TRUE(result.has_value())
+            << "Failed with " << result.error().code << ": " << result.error().message;
+    EXPECT_NE(result.value(), nullptr);
+}
+
+TEST(PreparedModelTest, configureExecutionBurstError) {
+    // setup test
+    const auto mockPreparedModel = MockPreparedModel::create();
+    const auto preparedModel =
+            PreparedModel::create(mockPreparedModel, /*executeSynchronously=*/true).value();
+    EXPECT_CALL(*mockPreparedModel, configureExecutionBurst(_, _, _, _))
+            .Times(1)
+            .WillOnce(
+                    makeConfigureExecutionBurstReturn(V1_0::ErrorStatus::GENERAL_FAILURE, nullptr));
+
+    // run test
+    const auto result = preparedModel->configureExecutionBurst();
+
+    // verify result
+    ASSERT_FALSE(result.has_value());
+    EXPECT_EQ(result.error().code, nn::ErrorStatus::GENERAL_FAILURE);
+}
+
+TEST(PreparedModelTest, configureExecutionBurstTransportFailure) {
+    // setup test
+    const auto mockPreparedModel = MockPreparedModel::create();
+    const auto preparedModel =
+            PreparedModel::create(mockPreparedModel, /*executeSynchronously=*/true).value();
+    EXPECT_CALL(*mockPreparedModel, configureExecutionBurst(_, _, _, _))
+            .Times(1)
+            .WillOnce(InvokeWithoutArgs(makeGeneralTransportFailure));
+
+    // run test
+    const auto result = preparedModel->configureExecutionBurst();
+
+    // verify result
+    ASSERT_FALSE(result.has_value());
+    EXPECT_EQ(result.error().code, nn::ErrorStatus::GENERAL_FAILURE);
+}
+
+TEST(PreparedModelTest, configureExecutionBurstDeadObject) {
+    // setup test
+    const auto mockPreparedModel = MockPreparedModel::create();
+    const auto preparedModel =
+            PreparedModel::create(mockPreparedModel, /*executeSynchronously=*/true).value();
+    EXPECT_CALL(*mockPreparedModel, configureExecutionBurst(_, _, _, _))
+            .Times(1)
+            .WillOnce(InvokeWithoutArgs(makeDeadObjectFailure));
+
+    // run test
+    const auto result = preparedModel->configureExecutionBurst();
+
+    // verify result
+    ASSERT_FALSE(result.has_value());
+    EXPECT_EQ(result.error().code, nn::ErrorStatus::DEAD_OBJECT);
+}
 
 TEST(PreparedModelTest, getUnderlyingResource) {
     // setup test