From 58e1526e1fb8d97625e67b394c25c94e5bb00656 Mon Sep 17 00:00:00 2001
From: "wangwenhua1@huawei.com" <wangwenhua1@huawei.com>
Date: Mon, 28 Dec 2020 17:54:21 +0800
Subject: [PATCH] display model info

---
 tests/ut/ge/CMakeLists.txt                  |  1 +
 tests/ut/ge/common/model_helper_unittest.cc | 52 +++++++++++++++++++++
 tests/ut/ge/session/omg_omg_unittest.cc     |  6 +--
 3 files changed, 56 insertions(+), 3 deletions(-)
 create mode 100644 tests/ut/ge/common/model_helper_unittest.cc

diff --git a/tests/ut/ge/CMakeLists.txt b/tests/ut/ge/CMakeLists.txt
index 82db6876..f4fe22b1 100755
--- a/tests/ut/ge/CMakeLists.txt
+++ b/tests/ut/ge/CMakeLists.txt
@@ -678,6 +678,7 @@ set(MULTI_PARTS_TEST_FILES
     "common/format_transfer_fracz_nhwc_unittest.cc"
     "common/format_transfer_fracz_hwcn_unittest.cc"
     "common/ge_format_util_unittest.cc"
+    "common/model_helper_unittest.cc"
     "graph/variable_accelerate_ctrl_unittest.cc"
     "graph/build/logical_stream_allocator_unittest.cc"
     "graph/build/mem_assigner_unittest.cc"
diff --git a/tests/ut/ge/common/model_helper_unittest.cc b/tests/ut/ge/common/model_helper_unittest.cc
new file mode 100644
index 00000000..7ce8c2cf
--- /dev/null
+++ b/tests/ut/ge/common/model_helper_unittest.cc
@@ -0,0 +1,52 @@
+/**
+ * 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 "framework/common/helper/model_helper.h"
+#include "ge/model/ge_model.h"
+#undef private
+#undef protected
+
+#include "proto/task.pb.h"
+
+using namespace ge;
+using namespace std;
+
+
+class UtestModelHelper : public testing::Test {
+ protected:
+  void SetUp() override {}
+
+  void TearDown() override {}
+};
+
+TEST_F(UtestModelHelper, save_size_to_modeldef_failed)
+{
+  GeModelPtr ge_model = ge::MakeShared<ge::GeModel>();
+  ModelHelper model_helper;
+  EXPECT_EQ(ACL_ERROR_GE_MEMORY_ALLOCATION, model_helper.SaveSizeToModelDef(ge_model));
+}
+
+TEST_F(UtestModelHelper, save_size_to_modeldef)
+{
+  GeModelPtr ge_model = ge::MakeShared<ge::GeModel>();
+  std::shared_ptr<domi::ModelTaskDef> task = ge::MakeShared<domi::ModelTaskDef>();
+  ge_model->SetModelTaskDef(task);
+  ModelHelper model_helper;
+  EXPECT_EQ(SUCCESS, model_helper.SaveSizeToModelDef(ge_model));
+}
diff --git a/tests/ut/ge/session/omg_omg_unittest.cc b/tests/ut/ge/session/omg_omg_unittest.cc
index 918be70b..6516dd33 100644
--- a/tests/ut/ge/session/omg_omg_unittest.cc
+++ b/tests/ut/ge/session/omg_omg_unittest.cc
@@ -24,20 +24,20 @@
 using namespace ge;
 using namespace std;
 
-class UTEST_omg_omg : public testing::Test {
+class UtestOmg : public testing::Test {
  protected:
   void SetUp() override {}
 
   void TearDown() override {}
 };
 
-TEST_F(UTEST_omg_omg, display_model_info_failed)
+TEST_F(UtestOmg, display_model_info_failed)
 {
   ge::proto::ModelDef model_def;
   PrintModelInfo(&model_def);
 }
 
-TEST_F(UTEST_omg_omg, display_model_info_success)
+TEST_F(UtestOmg, display_model_info_success)
 {
   ge::proto::ModelDef model_def;
   auto attrs = model_def.mutable_attr();