diff --git a/tests/ut/ge/graph/load/davinci_model_unittest.cc b/tests/ut/ge/graph/load/davinci_model_unittest.cc index 15dba075..41453fc9 100644 --- a/tests/ut/ge/graph/load/davinci_model_unittest.cc +++ b/tests/ut/ge/graph/load/davinci_model_unittest.cc @@ -35,11 +35,14 @@ class UtestDavinciModel : public testing::Test { TEST_F(UtestDavinciModel, init_data_op) { DavinciModel model(0, nullptr); + model.ge_model_ = make_shared(); + model.runtime_param_.mem_base = (uint8_t *)0x08000000; + model.runtime_param_.mem_size = 5120000; ComputeGraphPtr graph = make_shared("default"); OpDescPtr op_input = CreateOpDesc("data", DATA); - AttrUtils::SetInt(op_input, ATTR_NAME_PARENT_NODE_INDEX, 0); GeTensorDesc tensor(GeShape(), FORMAT_NCHW, DT_FLOAT); + TensorUtils::SetSize(tensor, 512); op_input->AddInputDesc(tensor); op_input->AddOutputDesc(tensor); op_input->SetInputOffset({1024}); @@ -49,20 +52,24 @@ TEST_F(UtestDavinciModel, init_data_op) { OpDescPtr op_output = CreateOpDesc("output", NETOUTPUT); op_output->AddInputDesc(tensor); op_output->SetInputOffset({1024}); + op_output->SetSrcName( { "data" } ); + op_output->SetSrcIndex( { 0 } ); 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); + EXPECT_EQ(model.op_list_.size(), 2); } TEST_F(UtestDavinciModel, init_data_op_subgraph) { DavinciModel model(0, nullptr); + model.runtime_param_.mem_base = (uint8_t *)0x08000000; + model.runtime_param_.mem_size = 5120000; ComputeGraphPtr graph = make_shared("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); @@ -82,12 +89,17 @@ TEST_F(UtestDavinciModel, init_data_op_subgraph) { TEST_F(UtestDavinciModel, init_netoutput_op_subgraph) { DavinciModel model(0, nullptr); + model.runtime_param_.mem_base = (uint8_t *)0x08000000; + model.runtime_param_.mem_size = 5120000; ComputeGraphPtr graph = make_shared("default"); OpDescPtr op_output = CreateOpDesc("output", NETOUTPUT); + GeTensorDesc tensor(GeShape(), FORMAT_NCHW, DT_FLOAT); op_output->AddInputDesc(tensor); op_output->SetInputOffset({1024}); - NodePtr node_output = graph->AddNode(op_output); + op_output->SetSrcName( { "data" } ); + op_output->SetSrcIndex( { 0 } ); + NodePtr output = graph->AddNode(op_output); std::vector output_op_list; EXPECT_EQ(model.InitNetOutput(nullptr, node, output_op_list), SUCCESS);