1 /* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
2 Licensed under the Apache License, Version 2.0 (the "License");
3 you may not use this file except in compliance with the License.
4 You may obtain a copy of the License at
5 http://www.apache.org/licenses/LICENSE-2.0
6 Unless required by applicable law or agreed to in writing, software
7 distributed under the License is distributed on an "AS IS" BASIS,
8 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
9 See the License for the specific language governing permissions and
10 limitations under the License.
11 ==============================================================================*/
12 #include "tensorflow/cc/saved_model/util.h"
13
14 #include <string>
15
16 #include "tensorflow/core/lib/core/status_test_util.h"
17 #include "tensorflow/core/lib/io/path.h"
18 #include "tensorflow/core/platform/env.h"
19 #include "tensorflow/core/platform/test.h"
20 #include "tensorflow/core/protobuf/meta_graph.pb.h"
21 #include "tensorflow/core/protobuf/saved_object_graph.pb.h"
22
23 namespace tensorflow {
24 namespace saved_model {
25 namespace {
26
TEST(UtilTest,TestGetWriteVersionV2)27 TEST(UtilTest, TestGetWriteVersionV2) {
28 SavedModel saved_model_proto;
29 MetaGraphDef* meta_graphdef = saved_model_proto.add_meta_graphs();
30 auto* object_graph_def = meta_graphdef->mutable_object_graph_def();
31 object_graph_def->add_nodes();
32 EXPECT_EQ(GetWriteVersion(saved_model_proto), "2");
33 }
34
TEST(UtilTest,TestGetWriteVersionV1)35 TEST(UtilTest, TestGetWriteVersionV1) {
36 SavedModel saved_model_proto;
37 saved_model_proto.add_meta_graphs();
38 EXPECT_EQ(GetWriteVersion(saved_model_proto), "1");
39 saved_model_proto.add_meta_graphs();
40 EXPECT_EQ(GetWriteVersion(saved_model_proto), "1");
41 }
42
43 } // namespace
44 } // namespace saved_model
45 } // namespace tensorflow
46