1 /* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 #ifndef TENSORFLOW_CORE_DATA_SERVICE_TEST_UTIL_H_
16 #define TENSORFLOW_CORE_DATA_SERVICE_TEST_UTIL_H_
17
18 #include <functional>
19 #include <ostream>
20 #include <string>
21 #include <vector>
22
23 #include "tensorflow/core/data/service/common.pb.h"
24 #include "tensorflow/core/platform/protobuf.h"
25 #include "tensorflow/core/platform/statusor.h"
26 #include "tensorflow/core/platform/test.h"
27 #include "tensorflow/core/platform/tstring.h"
28 #include "tensorflow/core/platform/types.h"
29
30 namespace tensorflow {
31 namespace data {
32 namespace testing {
33
34 // Returns a test dataset representing
35 // tf.data.Dataset.range(range). Useful for testing dataset graph execution.
36 DatasetDef RangeDataset(int64_t range);
37
38 // Returns a test dataset representing
39 // tf.data.Dataset.range(range).map(lambda x: x*x).
40 DatasetDef RangeSquareDataset(int64_t range);
41
42 // Returns a test dataset representing
43 // tf.data.Dataset.range(range).shard(SHARD_HINT, SHARD_HINT).
44 DatasetDef RangeDatasetWithShardHint(int64_t range);
45
46 // Returns a test dataset representing
47 // tf.data.Dataset.range(100000000).repeat().
48 DatasetDef InfiniteDataset();
49
50 // Returns a test dataset representing
51 // tf.data.Dataset.from_tensor_slices(["filenames"]).interleave(
52 // lambda filepath: tf.data.TextLineDataset(filepath),
53 // cycle_length=10)
54 StatusOr<DatasetDef> InterleaveTextlineDataset(
55 const std::vector<tstring>& filenames,
56 const std::vector<tstring>& contents);
57
58 // Repeatedly calls `f()`, blocking until `f()` returns `false`.
59 //
60 // Returns an error if `f()` returns an error.
61 Status WaitWhile(std::function<StatusOr<bool>()> f);
62
63 // TODO(b/229726259): Make EqualsProto available in Googletest
64 // (Public feature request: https://github.com/google/googletest/issues/1761).
65 class ProtoStringMatcher {
66 public:
ProtoStringMatcher(const tensorflow::protobuf::Message & expected)67 explicit ProtoStringMatcher(const tensorflow::protobuf::Message& expected)
68 : expected_(expected.ShortDebugString()) {}
69
70 template <typename Message>
MatchAndExplain(const Message & p,::testing::MatchResultListener *)71 bool MatchAndExplain(const Message& p,
72 ::testing::MatchResultListener*) const {
73 return p.ShortDebugString() == expected_;
74 }
75
DescribeTo(::std::ostream * os)76 void DescribeTo(::std::ostream* os) const { *os << expected_; }
DescribeNegationTo(::std::ostream * os)77 void DescribeNegationTo(::std::ostream* os) const {
78 *os << "not equal to expected message: " << expected_;
79 }
80
81 private:
82 const std::string expected_;
83 };
84
EqualsProto(const tensorflow::protobuf::Message & x)85 inline ::testing::PolymorphicMatcher<ProtoStringMatcher> EqualsProto(
86 const tensorflow::protobuf::Message& x) {
87 return ::testing::MakePolymorphicMatcher(ProtoStringMatcher(x));
88 }
89 } // namespace testing
90 } // namespace data
91 } // namespace tensorflow
92
93 #endif // TENSORFLOW_CORE_DATA_SERVICE_TEST_UTIL_H_
94