xref: /aosp_15_r20/external/federated-compute/fcp/client/engine/tf_wrapper_test.cc (revision 14675a029014e728ec732f129a32e299b2da0601)
1 /*
2  * Copyright 2021 Google LLC
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 #include "fcp/client/engine/tf_wrapper.h"
17 
18 #include "google/protobuf/any.pb.h"
19 #include "gtest/gtest.h"
20 #include "absl/status/status.h"
21 #include "absl/status/statusor.h"
22 #include "fcp/testing/testing.h"
23 
24 namespace fcp {
25 namespace client {
26 namespace engine {
27 namespace {
28 
29 using ::google::protobuf::Any;
30 using ::tensorflow::ConfigProto;
31 
TEST(TfWrapperInitializeConfigProtoTest,InvalidConfigProtoWrongTypeUrl)32 TEST(TfWrapperInitializeConfigProtoTest, InvalidConfigProtoWrongTypeUrl) {
33   // Create an Any with a valid value but invalid type URL.
34   ConfigProto config_proto;
35   config_proto.mutable_graph_options()->set_timeline_step(123);
36   Any packed_config_proto;
37   packed_config_proto.PackFrom(config_proto);
38   packed_config_proto.set_type_url("invalid");
39 
40   absl::StatusOr<ConfigProto> result =
41       TensorFlowWrapper::InitializeConfigProto(packed_config_proto);
42 
43   EXPECT_THAT(result, IsCode(INVALID_ARGUMENT));
44 }
45 
TEST(TfWrapperInitializeConfigProtoTest,InvalidConfigProtoEmptyTypeUrl)46 TEST(TfWrapperInitializeConfigProtoTest, InvalidConfigProtoEmptyTypeUrl) {
47   // Create an Any with a valid value but empty type URL.
48   ConfigProto config_proto;
49   config_proto.mutable_graph_options()->set_timeline_step(123);
50   Any packed_config_proto;
51   packed_config_proto.PackFrom(config_proto);
52   packed_config_proto.clear_type_url();
53 
54   absl::StatusOr<ConfigProto> result =
55       TensorFlowWrapper::InitializeConfigProto(packed_config_proto);
56 
57   EXPECT_THAT(result, IsCode(INVALID_ARGUMENT));
58 }
59 
TEST(TfWrapperInitializeConfigProtoTest,InvalidConfigProtoValue)60 TEST(TfWrapperInitializeConfigProtoTest, InvalidConfigProtoValue) {
61   // Set the correct type URL, but an unparseable value.
62   Any packed_config_proto;
63   packed_config_proto.PackFrom(ConfigProto());
64   packed_config_proto.set_value("nonparseable");
65 
66   absl::StatusOr<ConfigProto> result =
67       TensorFlowWrapper::InitializeConfigProto(packed_config_proto);
68 
69   EXPECT_THAT(result, IsCode(INVALID_ARGUMENT));
70 }
71 
TEST(TfWrapperInitializeConfigProtoTest,ValidNonEmptyConfigProtoValue)72 TEST(TfWrapperInitializeConfigProtoTest, ValidNonEmptyConfigProtoValue) {
73   // Create an Any containing a valid, non-empty ConfigProto.
74   ConfigProto config_proto;
75   config_proto.mutable_graph_options()->set_timeline_step(123);
76   Any packed_config_proto;
77   packed_config_proto.PackFrom(config_proto);
78 
79   absl::StatusOr<ConfigProto> result =
80       TensorFlowWrapper::InitializeConfigProto(packed_config_proto);
81 
82   // A non-empty ConfigProto was provided, so it should be used as-is.
83   ASSERT_OK(result);
84   EXPECT_THAT(*result, EqualsProto(config_proto));
85 }
86 
TEST(TfWrapperInitializeConfigProtoTest,ValidEmptyConfigProtoValue)87 TEST(TfWrapperInitializeConfigProtoTest, ValidEmptyConfigProtoValue) {
88   // Create an Any containing an empty ConfigProto.
89   Any packed_config_proto;
90   packed_config_proto.PackFrom(ConfigProto());
91 
92   absl::StatusOr<ConfigProto> result =
93       TensorFlowWrapper::InitializeConfigProto(packed_config_proto);
94 
95   // No external ConfigProto was provided, so the hardcoded defaults should be
96   // used.
97   ConfigProto expected_config_proto;
98   expected_config_proto.mutable_graph_options()->set_place_pruned_graph(true);
99   expected_config_proto.mutable_experimental()
100       ->set_disable_output_partition_graphs(true);
101   expected_config_proto.mutable_experimental()->set_optimize_for_static_graph(
102       true);
103 
104   ASSERT_OK(result);
105   EXPECT_THAT(*result, EqualsProto(expected_config_proto));
106 }
107 
TEST(TfWrapperInitializeConfigProtoTest,ValidEmptyPackedConfigProtoValue)108 TEST(TfWrapperInitializeConfigProtoTest, ValidEmptyPackedConfigProtoValue) {
109   // Create an empty Any.
110   Any packed_config_proto;
111 
112   absl::StatusOr<ConfigProto> result =
113       TensorFlowWrapper::InitializeConfigProto(packed_config_proto);
114 
115   // No external ConfigProto was provided, so the hardcoded defaults should be
116   // used.
117   ConfigProto expected_config_proto;
118   expected_config_proto.mutable_graph_options()->set_place_pruned_graph(true);
119   expected_config_proto.mutable_experimental()
120       ->set_disable_output_partition_graphs(true);
121   expected_config_proto.mutable_experimental()->set_optimize_for_static_graph(
122       true);
123 
124   ASSERT_OK(result);
125   EXPECT_THAT(*result, EqualsProto(expected_config_proto));
126 }
127 
128 }  // namespace
129 }  // namespace engine
130 }  // namespace client
131 }  // namespace fcp
132