tflite: mtk_neuron: Add PrepareSyncDriver helper in neuron_delegate_test

This CL extract the logic of preparing SyncDriver, making the code looks
cleaner. This CL is expected to be no functional changes.

BUG=b:374245241
TEST=neuron_delegate_test on navi

Change-Id: Ia4b19b5f3a646ccb9fe0bf7f689dcfc64badcd64
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/tflite/+/6437796
Reviewed-by: Shik Chen <shik@chromium.org>
Tested-by: Tommy Chiang <ototot@google.com>
Commit-Queue: Tommy Chiang <ototot@google.com>
diff --git a/delegate/mtk_neuron/neuron_delegate_test.cc b/delegate/mtk_neuron/neuron_delegate_test.cc
index 859cfbe..506c7a0 100644
--- a/delegate/mtk_neuron/neuron_delegate_test.cc
+++ b/delegate/mtk_neuron/neuron_delegate_test.cc
@@ -254,23 +254,39 @@
   return buffer;
 }
 
+std::unique_ptr<SyncDriver> PrepareSyncDriver(
+    const std::unique_ptr<FlatBufferModel>& model, TfLiteDelegatePtr delegate,
+    const std::vector<float>& input_buffer) {
+  if (model == nullptr || delegate == nullptr) {
+    return nullptr;
+  }
+  auto driver = SyncDriver::Create(
+      std::move(delegate), FlatBufferModel::BuildFromModel(model->GetModel()));
+  if (driver == nullptr) {
+    return nullptr;
+  }
+  driver->SetInputTensorBuffer(
+      "input",
+      AllocateTfLiteAlignedTensorBuffer(input_buffer.size() * sizeof(float)));
+  if (driver->SetInputTensorData("input", input_buffer) != kTfLiteOk) {
+    return nullptr;
+  }
+  if (driver->AllocateBuffers() != kTfLiteOk) {
+    return nullptr;
+  }
+  return driver;
+}
+
 TEST(NeuronDelegateTest, InferenceSelfAborted) {
   TFLiteSettingsT settings = GetDefaultSettings();
 
   auto [model, buffer_size] = GetTimeConsumingModelAndBufferSize();
-  ASSERT_NE(model, nullptr);
   std::vector<float> input = GetRandomizedFloatData(buffer_size);
 
   // Get the expected output with the given input.
-  TfLiteDelegatePtr delegate = GetDelegateFromSettings(&settings);
-  ASSERT_NE(delegate, nullptr);
-  auto driver = SyncDriver::Create(
-      std::move(delegate), FlatBufferModel::BuildFromModel(model->GetModel()));
+  auto driver =
+      PrepareSyncDriver(model, GetDelegateFromSettings(&settings), input);
   ASSERT_NE(driver, nullptr);
-  driver->SetInputTensorBuffer(
-      "input", AllocateTfLiteAlignedTensorBuffer(buffer_size * sizeof(float)));
-  ASSERT_EQ(driver->SetInputTensorData("input", input), kTfLiteOk);
-  ASSERT_EQ(driver->AllocateBuffers(), kTfLiteOk);
   ASSERT_EQ(driver->Invoke(), kTfLiteOk);
   std::vector<float> expected_output =
       driver->GetOutputTensorData<float>("output");
@@ -279,15 +295,8 @@
   settings.mtk_neuron_settings->inference_abort_time_ms = 1;
 
   // Try inference the model with an abort time.
-  delegate = GetDelegateFromSettings(&settings);
-  ASSERT_NE(delegate, nullptr);
-  driver = SyncDriver::Create(
-      std::move(delegate), FlatBufferModel::BuildFromModel(model->GetModel()));
+  driver = PrepareSyncDriver(model, GetDelegateFromSettings(&settings), input);
   ASSERT_NE(driver, nullptr);
-  driver->SetInputTensorBuffer(
-      "input", AllocateTfLiteAlignedTensorBuffer(buffer_size * sizeof(float)));
-  ASSERT_EQ(driver->SetInputTensorData("input", input), kTfLiteOk);
-  ASSERT_EQ(driver->AllocateBuffers(), kTfLiteOk);
 
   // Check the model inference is properly aborted.
   absl::Time before_invoke = absl::Now();
@@ -309,30 +318,17 @@
   TFLiteSettingsT settings = GetDefaultSettings();
 
   auto [model, buffer_size] = GetTimeConsumingModelAndBufferSize();
-  ASSERT_NE(model, nullptr);
   std::vector<float> input = GetRandomizedFloatData(buffer_size);
 
   // Create the first driver without inferencing.
-  TfLiteDelegatePtr delegate1 = GetDelegateFromSettings(&settings);
-  ASSERT_NE(delegate1, nullptr);
-  auto driver1 = SyncDriver::Create(
-      std::move(delegate1), FlatBufferModel::BuildFromModel(model->GetModel()));
+  auto driver1 =
+      PrepareSyncDriver(model, GetDelegateFromSettings(&settings), input);
   ASSERT_NE(driver1, nullptr);
-  driver1->SetInputTensorBuffer(
-      "input", AllocateTfLiteAlignedTensorBuffer(buffer_size * sizeof(float)));
-  ASSERT_EQ(driver1->SetInputTensorData("input", input), kTfLiteOk);
-  ASSERT_EQ(driver1->AllocateBuffers(), kTfLiteOk);
 
   // Create the second driver without inferencing.
-  TfLiteDelegatePtr delegate2 = GetDelegateFromSettings(&settings);
-  ASSERT_NE(delegate2, nullptr);
-  auto driver2 = SyncDriver::Create(
-      std::move(delegate2), FlatBufferModel::BuildFromModel(model->GetModel()));
+  auto driver2 =
+      PrepareSyncDriver(model, GetDelegateFromSettings(&settings), input);
   ASSERT_NE(driver2, nullptr);
-  driver2->SetInputTensorBuffer(
-      "input", AllocateTfLiteAlignedTensorBuffer(buffer_size * sizeof(float)));
-  ASSERT_EQ(driver2->SetInputTensorData("input", input), kTfLiteOk);
-  ASSERT_EQ(driver2->AllocateBuffers(), kTfLiteOk);
 
   // Create a child thread to inference the first driver.
   absl::Duration time_used1;