1 /* Copyright 2016 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 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 23 namespace tensorflow { 24 namespace io { 25 26 static std::vector<int> InputBufferSizes() { 27 return {10, 100, 200, 500, 1000, 10000}; 28 } 29 30 static std::vector<int> OutputBufferSizes() { return {100, 200, 500, 1000}; } 31 32 static std::vector<int> NumCopies() { return {1, 50, 500}; } 33 34 static 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 51 static 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 59 typedef io::ZlibCompressionOptions CompressionOptions; 60 61 void 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 96 TEST(ZlibBuffers, DefaultOptions) { 97 TestAllCombinations(CompressionOptions::DEFAULT(), 98 CompressionOptions::DEFAULT()); 99 } 100 101 TEST(ZlibBuffers, RawDeflate) { 102 TestAllCombinations(CompressionOptions::RAW(), CompressionOptions::RAW()); 103 } 104 105 TEST(ZlibBuffers, Gzip) { 106 TestAllCombinations(CompressionOptions::GZIP(), CompressionOptions::GZIP()); 107 } 108 109 void 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 153 TEST(ZlibBuffers, MultipleWritesWithoutFlush) { 154 TestMultipleWrites(200, 200, 10); 155 } 156 157 TEST(ZlibBuffers, MultipleWriteCallsWithFlush) { 158 TestMultipleWrites(200, 200, 10, true); 159 } 160 161 TEST(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 195 void 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 212 void 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 256 void 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 295 TEST(ZlibInputStream, TellDefaultOptions) { 296 TestTell(CompressionOptions::DEFAULT(), CompressionOptions::DEFAULT()); 297 } 298 299 TEST(ZlibInputStream, TellRawDeflate) { 300 TestTell(CompressionOptions::RAW(), CompressionOptions::RAW()); 301 } 302 303 TEST(ZlibInputStream, TellGzip) { 304 TestTell(CompressionOptions::GZIP(), CompressionOptions::GZIP()); 305 } 306 307 TEST(ZlibInputStream, SkipNBytesDefaultOptions) { 308 TestSkipNBytes(CompressionOptions::DEFAULT(), CompressionOptions::DEFAULT()); 309 } 310 311 TEST(ZlibInputStream, SkipNBytesRawDeflate) { 312 TestSkipNBytes(CompressionOptions::RAW(), CompressionOptions::RAW()); 313 } 314 315 TEST(ZlibInputStream, SkipNBytesGzip) { 316 TestSkipNBytes(CompressionOptions::GZIP(), CompressionOptions::GZIP()); 317 } 318 319 } // namespace io 320 } // namespace tensorflow 321
