1/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7    http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15
16#include "tensorflow/core/lib/core/status_test_util.h"
17#include "tensorflow/core/lib/io/random_inputstream.h"
18#include "tensorflow/core/lib/io/zlib_compression_options.h"
19#include "tensorflow/core/lib/io/zlib_inputstream.h"
20#include "tensorflow/core/lib/io/zlib_outputbuffer.h"
21#include "tensorflow/core/lib/strings/strcat.h"
22
23namespace tensorflow {
24namespace io {
25
26static std::vector<int> InputBufferSizes() {
27  return {10, 100, 200, 500, 1000, 10000};
28}
29
30static std::vector<int> OutputBufferSizes() { return {100, 200, 500, 1000}; }
31
32static std::vector<int> NumCopies() { return {1, 50, 500}; }
33
34static string GetRecord() {
35  static const string lorem_ipsum =
36      "Lorem ipsum dolor sit amet, consectetur adipiscing elit."
37      " Fusce vehicula tincidunt libero sit amet ultrices. Vestibulum non "
38      "felis augue. Duis vitae augue id lectus lacinia congue et ut purus. "
39      "Donec auctor, nisl at dapibus volutpat, diam ante lacinia dolor, vel"
40      "dignissim lacus nisi sed purus. Duis fringilla nunc ac lacus sagittis"
41      " efficitur. Praesent tincidunt egestas eros, eu vehicula urna ultrices"
42      " et. Aliquam erat volutpat. Maecenas vehicula risus consequat risus"
43      " dictum, luctus tincidunt nibh imperdiet. Aenean bibendum ac erat"
44      " cursus scelerisque. Cras lacinia in enim dapibus iaculis. Nunc porta"
45      " felis lectus, ac tincidunt massa pharetra quis. Fusce feugiat dolor"
46      " vel ligula rutrum egestas. Donec vulputate quam eros, et commodo"
47      " purus lobortis sed.";
48  return lorem_ipsum;
49}
50
51static string GenTestString(int copies = 1) {
52  string result = "";
53  for (int i = 0; i < copies; i++) {
54    result += GetRecord();
55  }
56  return result;
57}
58
59typedef io::ZlibCompressionOptions CompressionOptions;
60
61void TestAllCombinations(CompressionOptions input_options,
62                         CompressionOptions output_options) {
63  Env* env = Env::Default();
64  string fname = testing::TmpDir() + "/zlib_buffers_test";
65  for (auto file_size : NumCopies()) {
66    // Write to compressed file
67    string data = GenTestString(file_size);
68    for (auto input_buf_size : InputBufferSizes()) {
69      for (auto output_buf_size : OutputBufferSizes()) {
70        std::unique_ptr<WritableFile> file_writer;
71        TF_ASSERT_OK(env->NewWritableFile(fname, &file_writer));
72        string result;
73
74        ZlibOutputBuffer out(file_writer.get(), input_buf_size, output_buf_size,
75                             output_options);
76        TF_ASSERT_OK(out.Init());
77
78        TF_ASSERT_OK(out.Append(StringPiece(data)));
79        TF_ASSERT_OK(out.Close());
80        TF_ASSERT_OK(file_writer->Flush());
81        TF_ASSERT_OK(file_writer->Close());
82
83        std::unique_ptr<RandomAccessFile> file_reader;
84        TF_ASSERT_OK(env->NewRandomAccessFile(fname, &file_reader));
85        std::unique_ptr<RandomAccessInputStream> input_stream(
86            new RandomAccessInputStream(file_reader.get()));
87        ZlibInputStream in(input_stream.get(), input_buf_size, output_buf_size,
88                           input_options);
89        TF_ASSERT_OK(in.ReadNBytes(data.size(), &result));
90        EXPECT_EQ(result, data);
91      }
92    }
93  }
94}
95
96TEST(ZlibBuffers, DefaultOptions) {
97  TestAllCombinations(CompressionOptions::DEFAULT(),
98                      CompressionOptions::DEFAULT());
99}
100
101TEST(ZlibBuffers, RawDeflate) {
102  TestAllCombinations(CompressionOptions::RAW(), CompressionOptions::RAW());
103}
104
105TEST(ZlibBuffers, Gzip) {
106  TestAllCombinations(CompressionOptions::GZIP(), CompressionOptions::GZIP());
107}
108
109void TestMultipleWrites(uint8 input_buf_size, uint8 output_buf_size,
110                        int num_writes, bool with_flush = false) {
111  Env* env = Env::Default();
112  CompressionOptions input_options = CompressionOptions::DEFAULT();
113  CompressionOptions output_options = CompressionOptions::DEFAULT();
114
115  string fname = testing::TmpDir() + "/zlib_buffers_test";
116  string data = GenTestString();
117  std::unique_ptr<WritableFile> file_writer;
118  string actual_result;
119  string expected_result;
120
121  TF_ASSERT_OK(env->NewWritableFile(fname, &file_writer));
122  ZlibOutputBuffer out(file_writer.get(), input_buf_size, output_buf_size,
123                       output_options);
124  TF_ASSERT_OK(out.Init());
125
126  for (int i = 0; i < num_writes; i++) {
127    TF_ASSERT_OK(out.Append(StringPiece(data)));
128    if (with_flush) {
129      TF_ASSERT_OK(out.Flush());
130    }
131    strings::StrAppend(&expected_result, data);
132  }
133  TF_ASSERT_OK(out.Close());
134  TF_ASSERT_OK(file_writer->Flush());
135  TF_ASSERT_OK(file_writer->Close());
136
137  std::unique_ptr<RandomAccessFile> file_reader;
138  TF_ASSERT_OK(env->NewRandomAccessFile(fname, &file_reader));
139  std::unique_ptr<RandomAccessInputStream> input_stream(
140      new RandomAccessInputStream(file_reader.get()));
141  ZlibInputStream in(input_stream.get(), input_buf_size, output_buf_size,
142                     input_options);
143
144  for (int i = 0; i < num_writes; i++) {
145    string decompressed_output;
146    TF_ASSERT_OK(in.ReadNBytes(data.size(), &decompressed_output));
147    strings::StrAppend(&actual_result, decompressed_output);
148  }
149
150  EXPECT_EQ(actual_result, expected_result);
151}
152
153TEST(ZlibBuffers, MultipleWritesWithoutFlush) {
154  TestMultipleWrites(200, 200, 10);
155}
156
157TEST(ZlibBuffers, MultipleWriteCallsWithFlush) {
158  TestMultipleWrites(200, 200, 10, true);
159}
160
161TEST(ZlibInputStream, FailsToReadIfWindowBitsAreIncompatible) {
162  Env* env = Env::Default();
163  string fname = testing::TmpDir() + "/zlib_buffers_test";
164  CompressionOptions output_options = CompressionOptions::DEFAULT();
165  CompressionOptions input_options = CompressionOptions::DEFAULT();
166  int input_buf_size = 200, output_buf_size = 200;
167  output_options.window_bits = MAX_WBITS;
168  // inflate() has smaller history buffer.
169  input_options.window_bits = output_options.window_bits - 1;
170
171  string data = GenTestString(10);
172  std::unique_ptr<WritableFile> file_writer;
173  TF_ASSERT_OK(env->NewWritableFile(fname, &file_writer));
174  string result;
175  ZlibOutputBuffer out(file_writer.get(), input_buf_size, output_buf_size,
176                       output_options);
177  TF_ASSERT_OK(out.Init());
178
179  TF_ASSERT_OK(out.Append(StringPiece(data)));
180  TF_ASSERT_OK(out.Close());
181  TF_ASSERT_OK(file_writer->Flush());
182  TF_ASSERT_OK(file_writer->Close());
183
184  std::unique_ptr<RandomAccessFile> file_reader;
185  TF_ASSERT_OK(env->NewRandomAccessFile(fname, &file_reader));
186  std::unique_ptr<RandomAccessInputStream> input_stream(
187      new RandomAccessInputStream(file_reader.get()));
188  ZlibInputStream in(input_stream.get(), input_buf_size, output_buf_size,
189                     input_options);
190  Status read_status = in.ReadNBytes(data.size(), &result);
191  CHECK_EQ(read_status.code(), error::DATA_LOSS);
192  CHECK(read_status.error_message().find("inflate() failed") != string::npos);
193}
194
195void WriteCompressedFile(Env* env, const string& fname, int input_buf_size,
196                         int output_buf_size,
197                         const CompressionOptions& output_options,
198                         const string& data) {
199  std::unique_ptr<WritableFile> file_writer;
200  TF_ASSERT_OK(env->NewWritableFile(fname, &file_writer));
201
202  ZlibOutputBuffer out(file_writer.get(), input_buf_size, output_buf_size,
203                       output_options);
204  TF_ASSERT_OK(out.Init());
205
206  TF_ASSERT_OK(out.Append(StringPiece(data)));
207  TF_ASSERT_OK(out.Close());
208  TF_ASSERT_OK(file_writer->Flush());
209  TF_ASSERT_OK(file_writer->Close());
210}
211
212void TestTell(CompressionOptions input_options,
213              CompressionOptions output_options) {
214  Env* env = Env::Default();
215  string fname = testing::TmpDir() + "/zlib_buffers_test";
216  for (auto file_size : NumCopies()) {
217    string data = GenTestString(file_size);
218    for (auto input_buf_size : InputBufferSizes()) {
219      for (auto output_buf_size : OutputBufferSizes()) {
220        // Write the compressed file.
221        WriteCompressedFile(env, fname, input_buf_size, output_buf_size,
222                            output_options, data);
223
224        // Boiler-plate to set up ZlibInputStream.
225        std::unique_ptr<RandomAccessFile> file_reader;
226        TF_ASSERT_OK(env->NewRandomAccessFile(fname, &file_reader));
227        std::unique_ptr<RandomAccessInputStream> input_stream(
228            new RandomAccessInputStream(file_reader.get()));
229        ZlibInputStream in(input_stream.get(), input_buf_size, output_buf_size,
230                           input_options);
231
232        string first_half(data, 0, data.size() / 2);
233        string bytes_read;
234
235        // Read the first half of the uncompressed file and expect that Tell()
236        // returns half the uncompressed length of the file.
237        TF_ASSERT_OK(in.ReadNBytes(first_half.size(), &bytes_read));
238        EXPECT_EQ(in.Tell(), first_half.size());
239        EXPECT_EQ(bytes_read, first_half);
240
241        // Read the remaining half of the uncompressed file and expect that
242        // Tell() points past the end of file.
243        string second_half;
244        TF_ASSERT_OK(
245            in.ReadNBytes(data.size() - first_half.size(), &second_half));
246        EXPECT_EQ(in.Tell(), data.size());
247        bytes_read.append(second_half);
248
249        // Expect that the file is correctly read.
250        EXPECT_EQ(bytes_read, data);
251      }
252    }
253  }
254}
255
256void TestSkipNBytes(CompressionOptions input_options,
257                    CompressionOptions output_options) {
258  Env* env = Env::Default();
259  string fname = testing::TmpDir() + "/zlib_buffers_test";
260  for (auto file_size : NumCopies()) {
261    string data = GenTestString(file_size);
262    for (auto input_buf_size : InputBufferSizes()) {
263      for (auto output_buf_size : OutputBufferSizes()) {
264        // Write the compressed file.
265        WriteCompressedFile(env, fname, input_buf_size, output_buf_size,
266                            output_options, data);
267
268        // Boiler-plate to set up ZlibInputStream.
269        std::unique_ptr<RandomAccessFile> file_reader;
270        TF_ASSERT_OK(env->NewRandomAccessFile(fname, &file_reader));
271        std::unique_ptr<RandomAccessInputStream> input_stream(
272            new RandomAccessInputStream(file_reader.get()));
273        ZlibInputStream in(input_stream.get(), input_buf_size, output_buf_size,
274                           input_options);
275
276        size_t data_half_size = data.size() / 2;
277        string second_half(data, data_half_size, data.size() - data_half_size);
278
279        // Skip past the first half of the file and expect Tell() returns
280        // correctly.
281        TF_ASSERT_OK(in.SkipNBytes(data_half_size));
282        EXPECT_EQ(in.Tell(), data_half_size);
283
284        // Expect that second half is read correctly and Tell() returns past
285        // end of file after reading complete file.
286        string bytes_read;
287        TF_ASSERT_OK(in.ReadNBytes(second_half.size(), &bytes_read));
288        EXPECT_EQ(bytes_read, second_half);
289        EXPECT_EQ(in.Tell(), data.size());
290      }
291    }
292  }
293}
294
295TEST(ZlibInputStream, TellDefaultOptions) {
296  TestTell(CompressionOptions::DEFAULT(), CompressionOptions::DEFAULT());
297}
298
299TEST(ZlibInputStream, TellRawDeflate) {
300  TestTell(CompressionOptions::RAW(), CompressionOptions::RAW());
301}
302
303TEST(ZlibInputStream, TellGzip) {
304  TestTell(CompressionOptions::GZIP(), CompressionOptions::GZIP());
305}
306
307TEST(ZlibInputStream, SkipNBytesDefaultOptions) {
308  TestSkipNBytes(CompressionOptions::DEFAULT(), CompressionOptions::DEFAULT());
309}
310
311TEST(ZlibInputStream, SkipNBytesRawDeflate) {
312  TestSkipNBytes(CompressionOptions::RAW(), CompressionOptions::RAW());
313}
314
315TEST(ZlibInputStream, SkipNBytesGzip) {
316  TestSkipNBytes(CompressionOptions::GZIP(), CompressionOptions::GZIP());
317}
318
319}  // namespace io
320}  // namespace tensorflow
321