From caf1a427b8e0f952e5a4d5d8dfe1d3f43535a665 Mon Sep 17 00:00:00 2001
From: zhangxiaokun <zhang.xiaokun@huawei.com>
Date: Tue, 29 Dec 2020 20:12:18 +0800
Subject: [PATCH] Add UT

---
 tests/ut/ge/CMakeLists.txt                    |   1 +
 .../ge/graph/load/davinci_model_unittest.cc   | 100 ++++++++++++++++++
 2 files changed, 101 insertions(+)
 create mode 100644 tests/ut/ge/graph/load/davinci_model_unittest.cc

diff --git a/tests/ut/ge/CMakeLists.txt b/tests/ut/ge/CMakeLists.txt
index 175774bb..ebaf7708 100755
--- a/tests/ut/ge/CMakeLists.txt
+++ b/tests/ut/ge/CMakeLists.txt
@@ -565,6 +565,7 @@ set(DISTINCT_GRAPH_LOAD_TEST_FILES
     "graph/load/end_graph_task_unittest.cc"
     "graph/load/new_model_manager_event_manager_unittest.cc"
     #"graph/load/output_net_output_unittest.cc"
+    "graph/load/davinci_model_unittest.cc"
     "graph/load/tbe_handle_store_unittest.cc"
     "graph/load/hccl_task_info_unittest.cc"
     "graph/load/kernel_ex_task_info_unittest.cc"
diff --git a/tests/ut/ge/graph/load/davinci_model_unittest.cc b/tests/ut/ge/graph/load/davinci_model_unittest.cc
new file mode 100644
index 00000000..15dba075
--- /dev/null
+++ b/tests/ut/ge/graph/load/davinci_model_unittest.cc
@@ -0,0 +1,100 @@
+/**
+ * Copyright 2019-2020 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>
+
+#define private public
+#define protected public
+#include "graph/load/new_model_manager/davinci_model.h"
+
+using namespace std;
+
+namespace ge {
+extern OpDescPtr CreateOpDesc(string name, string type);
+
+class UtestDavinciModel : public testing::Test {
+ protected:
+  void SetUp() {}
+
+  void TearDown() {}
+};
+
+
+TEST_F(UtestDavinciModel, init_data_op) {
+  DavinciModel model(0, nullptr);
+  ComputeGraphPtr graph = make_shared<ComputeGraph>("default");
+
+  OpDescPtr op_input = CreateOpDesc("data", DATA);
+  AttrUtils::SetInt(op_input, ATTR_NAME_PARENT_NODE_INDEX, 0);
+  GeTensorDesc tensor(GeShape(), FORMAT_NCHW, DT_FLOAT);
+  op_input->AddInputDesc(tensor);
+  op_input->AddOutputDesc(tensor);
+  op_input->SetInputOffset({1024});
+  op_input->SetOutputOffset({5120});
+  NodePtr node_input = graph->AddNode(op_input);
+
+  OpDescPtr op_output = CreateOpDesc("output", NETOUTPUT);
+  op_output->AddInputDesc(tensor);
+  op_output->SetInputOffset({1024});
+  NodePtr node_output = graph->AddNode(op_output);
+
+  EXPECT_EQ(model.InitNodes(graph), SUCCESS);
+
+  EXPECT_EQ(model.input_addrs_list_.size(), 1);
+  EXPECT_EQ(model.output_addrs_list_.size(), 1);
+}
+
+TEST_F(UtestDavinciModel, init_data_op_subgraph) {
+  DavinciModel model(0, nullptr);
+  ComputeGraphPtr graph = make_shared<ComputeGraph>("default");
+
+  OpDescPtr op_input = CreateOpDesc("data", DATA);
+  AttrUtils::SetInt(op_input, ATTR_NAME_PARENT_NODE_INDEX, 0);
+  GeTensorDesc tensor(GeShape(), FORMAT_NCHW, DT_FLOAT);
+  op_input->AddInputDesc(tensor);
+  op_input->AddOutputDesc(tensor);
+  op_input->SetInputOffset({1024});
+  op_input->SetOutputOffset({5120});
+  NodePtr node = graph->AddNode(op_input);
+
+  uint32_t data_op_index = 0;
+  map<uint32_t, OpDescPtr> data_by_index;
+  EXPECT_EQ(model.InitDataOp(nullptr, node, data_op_index, data_by_index), SUCCESS);
+
+  EXPECT_EQ(model.input_addrs_list_.size(), 0);
+  EXPECT_EQ(model.output_addrs_list_.size(), 0);
+  EXPECT_EQ(data_op_index, 0);
+  EXPECT_TRUE(data_by_index.empty());
+}
+
+TEST_F(UtestDavinciModel, init_netoutput_op_subgraph) {
+  DavinciModel model(0, nullptr);
+  ComputeGraphPtr graph = make_shared<ComputeGraph>("default");
+
+  OpDescPtr op_output = CreateOpDesc("output", NETOUTPUT);
+  op_output->AddInputDesc(tensor);
+  op_output->SetInputOffset({1024});
+  NodePtr node_output = graph->AddNode(op_output);
+
+  std::vector<OpDescPtr> output_op_list;
+  EXPECT_EQ(model.InitNetOutput(nullptr, node, output_op_list), SUCCESS);
+
+  EXPECT_EQ(model.input_addrs_list_.size(), 0);
+  EXPECT_EQ(model.output_addrs_list_.size(), 0);
+  EXPECT_TRUE(output_op_list.empty());
+}
+
+}  // namespace ge