From 241b2166ced7fb0fa661b0e34f688ae560940b81 Mon Sep 17 00:00:00 2001
From: chuxing <chuxing@huawei.com>
Date: Mon, 22 Mar 2021 15:44:45 +0800
Subject: [PATCH] add ut

---
 .../compiledsubgraph/known_node_executor.cc   |  6 +-
 .../compiledsubgraph/known_node_executor.h    |  3 +
 tests/ut/ge/CMakeLists.txt                    |  1 +
 .../ge/hybrid/known_node_executor_unittest.cc | 65 +++++++++++++++++++
 4 files changed, 74 insertions(+), 1 deletion(-)
 create mode 100644 tests/ut/ge/hybrid/known_node_executor_unittest.cc

diff --git a/ge/hybrid/node_executor/compiledsubgraph/known_node_executor.cc b/ge/hybrid/node_executor/compiledsubgraph/known_node_executor.cc
index 918603f0..d79fcc86 100755
--- a/ge/hybrid/node_executor/compiledsubgraph/known_node_executor.cc
+++ b/ge/hybrid/node_executor/compiledsubgraph/known_node_executor.cc
@@ -146,10 +146,14 @@ Status KnownNodeTask::InitDavinciModel() {
   int32_t device_id = 0;
   GE_CHK_RT_RET(rtGetDevice(&device_id));
   davinci_model_->SetDeviceId(static_cast<uint32_t>(device_id));
-  GE_CHK_STATUS_RET(davinci_model_->Init(), "[Init][Model] Failed to init davinci model.");
+  GE_CHK_STATUS_RET(DoInitDavinciModel(), "[Init][Model] Failed to init davinci model.");
   return SUCCESS;
 }
 
+Status KnownNodeTask::DoInitDavinciModel() {
+  return davinci_model_->Init();
+}
+
 Status KnownNodeExecutor::PrepareTask(NodeTask &task, TaskContext &context) const {
   GELOGD("[%s] KnownNodeExecutor::PrepareTask in.", context.GetNodeName());
   RECORD_EXECUTION_EVENT(context.GetExecutionContext(), context.GetNodeName(), "[KnownNodeExecutorPrepareTask] Start");
diff --git a/ge/hybrid/node_executor/compiledsubgraph/known_node_executor.h b/ge/hybrid/node_executor/compiledsubgraph/known_node_executor.h
index 762e9991..5eed528a 100644
--- a/ge/hybrid/node_executor/compiledsubgraph/known_node_executor.h
+++ b/ge/hybrid/node_executor/compiledsubgraph/known_node_executor.h
@@ -37,6 +37,9 @@ class KnownNodeTask : public NodeTask {
   Status ExecuteAsync(TaskContext &context, std::function<void()> done_callback) override;
   Status Init(TaskContext &context) override;
   Status InitDavinciModel();
+
+ protected:
+  virtual Status DoInitDavinciModel();
  private:
   std::shared_ptr<DavinciModel> davinci_model_ = nullptr;
   bool load_flag_ = false;
diff --git a/tests/ut/ge/CMakeLists.txt b/tests/ut/ge/CMakeLists.txt
index 3caba788..4c4a72a3 100755
--- a/tests/ut/ge/CMakeLists.txt
+++ b/tests/ut/ge/CMakeLists.txt
@@ -793,6 +793,7 @@ set(PROFILING_MNG_TEST_FILES
 
 set(HYBRID_TEST_FILES
     "hybrid/ge_hybrid_unittest.cc"
+    "hybrid/known_node_executor_unittest.cc"
 )
 
 set(OTHERS_TEST_FILES
diff --git a/tests/ut/ge/hybrid/known_node_executor_unittest.cc b/tests/ut/ge/hybrid/known_node_executor_unittest.cc
new file mode 100644
index 00000000..21b3faa9
--- /dev/null
+++ b/tests/ut/ge/hybrid/known_node_executor_unittest.cc
@@ -0,0 +1,65 @@
+/**
+ * Copyright 2019-2021 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <gtest/gtest.h>
+#include <gmock/gmock.h>
+#include <vector>
+#include <memory>
+
+#define protected public
+#define private public
+#include "hybrid/node_executor/compiledsubgraph/known_node_executor.h"
+#undef private
+#undef protected
+#include "graph/manager/graph_mem_allocator.h"
+
+using namespace std;
+using namespace testing;
+using namespace ge;
+using namespace hybrid;
+
+class UnknownNodeExecutorTest : public testing::Test {
+ protected:
+  void SetUp() {}
+  void TearDown() {}
+};
+
+namespace {
+class KnownNodeTaskMock : public KnownNodeTask {
+ public:
+  KnownNodeTaskMock(std::shared_ptr<DavinciModel> davinci_model): KnownNodeTask(davinci_model) {};
+  ~KnownNodeTaskMock() override = default;
+  MOCK_METHOD0(DoInitDavinciModel, Status());
+};
+}
+
+TEST_F(UnknownNodeExecutorTest, test_init_davinci_model) {
+  auto davinci_model = std::make_shared<DavinciModel>(0, nullptr);
+  davinci_model->SetDeviceId(0);
+  davinci_model->SetKnownNode(true);
+
+  MemManager::Instance().Initialize(std::vector<rtMemType_t>({RT_MEMORY_HBM}));
+
+  auto ge_model = make_shared<GeModel>();
+  AttrUtils::SetInt(ge_model, ATTR_MODEL_VAR_SIZE, 0);
+  AttrUtils::SetInt(ge_model, ATTR_MODEL_MEMORY_SIZE, 1024);
+  davinci_model->Assign(ge_model);
+
+  KnownNodeTaskMock mock(davinci_model);
+  EXPECT_CALL(mock, DoInitDavinciModel).WillOnce(::testing::Return(SUCCESS));
+  ASSERT_EQ(mock.InitDavinciModel(), SUCCESS);
+  MemManager::Instance().Finalize();
+}
\ No newline at end of file