From 35beed3c7e64b4bb4b39a2553d73ca80102f7398 Mon Sep 17 00:00:00 2001
From: zhangxiaokun <zhang.xiaokun@huawei.com>
Date: Tue, 29 Dec 2020 22:08:09 +0800
Subject: [PATCH] Add UT

---
 .../ge/graph/load/davinci_model_unittest.cc   | 56 +++++++++++++++++++
 1 file changed, 56 insertions(+)

diff --git a/tests/ut/ge/graph/load/davinci_model_unittest.cc b/tests/ut/ge/graph/load/davinci_model_unittest.cc
index 8374e9dc..37516c6e 100644
--- a/tests/ut/ge/graph/load/davinci_model_unittest.cc
+++ b/tests/ut/ge/graph/load/davinci_model_unittest.cc
@@ -18,6 +18,7 @@
 
 #define private public
 #define protected public
+#include "graph/utils/graph_utils.h"
 #include "graph/load/new_model_manager/davinci_model.h"
 
 using namespace std;
@@ -32,6 +33,61 @@ class UtestDavinciModel : public testing::Test {
   void TearDown() {}
 };
 
+TEST_F(UtestDavinciModel, init_success) {
+  DavinciModel model(0, nullptr);
+  ComputeGraphPtr graph = make_shared<ComputeGraph>("default");
+
+  GeModelPtr ge_model = make_shared<GeModel>();
+  ge_model->SetGraph(GraphUtils::CreateGraphFromComputeGraph(graph));
+  AttrUtils::SetInt(ge_model, ATTR_NODEL_MEMORY_SIZE, 5120000);
+  AttrUtils::SetInt(ge_model, ATTR_NODEL_STREAM_NUM, 1);
+
+  shared_ptr<domi::ModelTaskDef> model_task_def = make_shared<domi::ModelTaskDef>();
+  ge_model->SetModelTaskDef(model_task_def);
+
+  GeTensorDesc tensor(GeShape(), FORMAT_NCHW, DT_FLOAT);
+  TensorUtils::SetSize(tensor, 512);
+
+  OpDescPtr op_input = CreateOpDesc("data", DATA);
+  op_input->AddInputDesc(tensor);
+  op_input->AddOutputDesc(tensor);
+  op_input->SetInputOffset({1024});
+  op_input->SetOutputOffset({1024});
+  NodePtr node_input = graph->AddNode(op_input);
+
+  OpDescPtr op_memcpy = CreateOpDesc("memcpy", MEMCPYASYNC);
+  op_memcpy->AddInputDesc(tensor);
+  op_memcpy->AddOutputDesc(tensor);
+  op_memcpy->SetInputOffset({1024});
+  op_memcpy->SetOutputOffset({5120});
+  NodePtr node_memcpy = graph->AddNode(op_memcpy);
+
+  OpDescPtr op_output = CreateOpDesc("output", NETOUTPUT);
+  op_output->AddInputDesc(tensor);
+  op_output->SetInputOffset({5120});
+  op_output->SetSrcName( { "memcpy" } );
+  op_output->SetSrcIndex( { 0 } );
+  NodePtr node_output = graph->AddNode(op_output);
+
+  domi::TaskDef *task_def = model_task_def->AddTask();
+  task_def->set_stream_id(0);
+  task_def->set_type(RT_MODEL_TASK_MEMCPY_ASYNC);
+  domi::MemcpyAsyncDef *memcpy_async = task_def->mutable_memcpy_async();
+  memcpy_async->set_src(1024);
+  memcpy_async->set_dst(5120);
+  memcpy_async->set_dst_max(512);
+  memcpy_async->set_count(1);
+  memcpy_async->set_kind(RT_MEMCPY_DEVICE_TO_DEVICE);
+  memcpy_async->set_op_index(1);
+
+  EXPECT_EQ(model.Assign(ge_model), SUCCESS);
+  EXPECT_EQ(model.Init(), SUCCESS);
+
+  EXPECT_EQ(model.input_addrs_list_.size(), 1);
+  EXPECT_EQ(model.output_addrs_list_.size(), 1);
+  EXPECT_EQ(model.task_list_.size(), 1);
+  EXPECT_EQ(model.op_list_.size(), 3);
+}
 
 TEST_F(UtestDavinciModel, init_data_op) {
   DavinciModel model(0, nullptr);