xref: /aosp_15_r20/external/tensorflow/tensorflow/core/lib/io/buffered_inputstream_test.cc (revision b6fb3261f9314811a0f4371741dbb8839866f948) 
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/io/buffered_inputstream.h"
17 
18 #include "tensorflow/core/lib/core/status_test_util.h"
19 #include "tensorflow/core/lib/io/random_inputstream.h"
20 #include "tensorflow/core/platform/env.h"
21 #include "tensorflow/core/platform/test.h"
22 #include "tensorflow/core/platform/test_benchmark.h"
23 
24 namespace tensorflow {
25 namespace io {
26 namespace {
27 
BufferSizes()28 static std::vector<int> BufferSizes() {
29   return {1,  2,  3,  4,  5,  6,  7,  8,  9,  10,   11,
30           12, 13, 14, 15, 16, 17, 18, 19, 20, 65536};
31 }
32 
33 // This class will only return OutOfRange error once to make sure that
34 // BufferedInputStream is able to cache the error.
35 class ReadOnceInputStream : public InputStreamInterface {
36  public:
ReadOnceInputStream()37   ReadOnceInputStream() : start_(true) {}
38 
ReadNBytes(int64_t bytes_to_read,tstring * result)39   virtual Status ReadNBytes(int64_t bytes_to_read, tstring* result) {
40     if (bytes_to_read < 11) {
41       return errors::InvalidArgument("Not reading all bytes: ", bytes_to_read);
42     }
43     if (start_) {
44       *result = "0123456789";
45       start_ = false;
46       return errors::OutOfRange("Out of range.");
47     }
48     return errors::InvalidArgument(
49         "Redudant call to ReadNBytes after an OutOfRange error.");
50   }
51 
Tell() const52   int64_t Tell() const override { return start_ ? 0 : 10; }
53 
54   // Resets the stream to the beginning.
Reset()55   Status Reset() override {
56     start_ = true;
57     return OkStatus();
58   }
59 
60  private:
61   bool start_;
62 };
63 
TEST(BufferedInputStream,ReadLine_Empty)64 TEST(BufferedInputStream, ReadLine_Empty) {
65   Env* env = Env::Default();
66   string fname;
67   ASSERT_TRUE(env->LocalTempFilename(&fname));
68   TF_ASSERT_OK(WriteStringToFile(env, fname, ""));
69   std::unique_ptr<RandomAccessFile> file;
70   TF_ASSERT_OK(env->NewRandomAccessFile(fname, &file));
71 
72   for (auto buf_size : BufferSizes()) {
73     std::unique_ptr<RandomAccessInputStream> input_stream(
74         new RandomAccessInputStream(file.get()));
75     BufferedInputStream in(input_stream.get(), buf_size);
76     string line;
77     EXPECT_TRUE(errors::IsOutOfRange(in.ReadLine(&line)));
78   }
79 }
80 
TEST(BufferedInputStream,ReadLine1)81 TEST(BufferedInputStream, ReadLine1) {
82   Env* env = Env::Default();
83   string fname;
84   ASSERT_TRUE(env->LocalTempFilename(&fname));
85   TF_ASSERT_OK(
86       WriteStringToFile(env, fname, "line one\nline two\nline three\n"));
87   std::unique_ptr<RandomAccessFile> file;
88   TF_ASSERT_OK(env->NewRandomAccessFile(fname, &file));
89 
90   for (auto buf_size : BufferSizes()) {
91     std::unique_ptr<RandomAccessInputStream> input_stream(
92         new RandomAccessInputStream(file.get()));
93     BufferedInputStream in(input_stream.get(), buf_size);
94     string line;
95     TF_ASSERT_OK(in.ReadLine(&line));
96     EXPECT_EQ(line, "line one");
97     TF_ASSERT_OK(in.ReadLine(&line));
98     EXPECT_EQ(line, "line two");
99     TF_ASSERT_OK(in.ReadLine(&line));
100     EXPECT_EQ(line, "line three");
101     EXPECT_TRUE(errors::IsOutOfRange(in.ReadLine(&line)));
102     // A second call should also return end of file
103     EXPECT_TRUE(errors::IsOutOfRange(in.ReadLine(&line)));
104   }
105 }
106 
TEST(BufferedInputStream,ReadLine_NoTrailingNewLine)107 TEST(BufferedInputStream, ReadLine_NoTrailingNewLine) {
108   Env* env = Env::Default();
109   string fname;
110   ASSERT_TRUE(env->LocalTempFilename(&fname));
111   TF_ASSERT_OK(WriteStringToFile(env, fname, "line one\nline two\nline three"));
112   std::unique_ptr<RandomAccessFile> file;
113   TF_ASSERT_OK(env->NewRandomAccessFile(fname, &file));
114 
115   for (auto buf_size : BufferSizes()) {
116     std::unique_ptr<RandomAccessInputStream> input_stream(
117         new RandomAccessInputStream(file.get()));
118     BufferedInputStream in(input_stream.get(), buf_size);
119     string line;
120     TF_ASSERT_OK(in.ReadLine(&line));
121     EXPECT_EQ(line, "line one");
122     TF_ASSERT_OK(in.ReadLine(&line));
123     EXPECT_EQ(line, "line two");
124     TF_ASSERT_OK(in.ReadLine(&line));
125     EXPECT_EQ(line, "line three");
126     EXPECT_TRUE(errors::IsOutOfRange(in.ReadLine(&line)));
127     // A second call should also return end of file
128     EXPECT_TRUE(errors::IsOutOfRange(in.ReadLine(&line)));
129   }
130 }
131 
TEST(BufferedInputStream,ReadLine_EmptyLines)132 TEST(BufferedInputStream, ReadLine_EmptyLines) {
133   Env* env = Env::Default();
134   string fname;
135   ASSERT_TRUE(env->LocalTempFilename(&fname));
136   TF_ASSERT_OK(
137       WriteStringToFile(env, fname, "line one\n\n\nline two\nline three"));
138   std::unique_ptr<RandomAccessFile> file;
139   TF_ASSERT_OK(env->NewRandomAccessFile(fname, &file));
140 
141   for (auto buf_size : BufferSizes()) {
142     std::unique_ptr<RandomAccessInputStream> input_stream(
143         new RandomAccessInputStream(file.get()));
144     BufferedInputStream in(input_stream.get(), buf_size);
145     string line;
146     TF_ASSERT_OK(in.ReadLine(&line));
147     EXPECT_EQ(line, "line one");
148     TF_ASSERT_OK(in.ReadLine(&line));
149     EXPECT_EQ(line, "");
150     TF_ASSERT_OK(in.ReadLine(&line));
151     EXPECT_EQ(line, "");
152     TF_ASSERT_OK(in.ReadLine(&line));
153     EXPECT_EQ(line, "line two");
154     TF_ASSERT_OK(in.ReadLine(&line));
155     EXPECT_EQ(line, "line three");
156     EXPECT_TRUE(errors::IsOutOfRange(in.ReadLine(&line)));
157     // A second call should also return end of file
158     EXPECT_TRUE(errors::IsOutOfRange(in.ReadLine(&line)));
159   }
160 }
161 
TEST(BufferedInputStream,ReadLine_CRLF)162 TEST(BufferedInputStream, ReadLine_CRLF) {
163   Env* env = Env::Default();
164   string fname;
165   ASSERT_TRUE(env->LocalTempFilename(&fname));
166   TF_ASSERT_OK(WriteStringToFile(env, fname,
167                                  "line one\r\n\r\n\r\nline two\r\nline three"));
168   std::unique_ptr<RandomAccessFile> file;
169   TF_ASSERT_OK(env->NewRandomAccessFile(fname, &file));
170 
171   for (auto buf_size : BufferSizes()) {
172     std::unique_ptr<RandomAccessInputStream> input_stream(
173         new RandomAccessInputStream(file.get()));
174     BufferedInputStream in(input_stream.get(), buf_size);
175     string line;
176     TF_ASSERT_OK(in.ReadLine(&line));
177     EXPECT_EQ(line, "line one");
178     TF_ASSERT_OK(in.ReadLine(&line));
179     EXPECT_EQ(line, "");
180     TF_ASSERT_OK(in.ReadLine(&line));
181     EXPECT_EQ(line, "");
182     TF_ASSERT_OK(in.ReadLine(&line));
183     EXPECT_EQ(line, "line two");
184     TF_ASSERT_OK(in.ReadLine(&line));
185     EXPECT_EQ(line, "line three");
186     EXPECT_TRUE(errors::IsOutOfRange(in.ReadLine(&line)));
187     // A second call should also return end of file
188     EXPECT_TRUE(errors::IsOutOfRange(in.ReadLine(&line)));
189   }
190 }
191 
TEST(BufferedInputStream,SkipLine1)192 TEST(BufferedInputStream, SkipLine1) {
193   Env* env = Env::Default();
194   string fname;
195   ASSERT_TRUE(env->LocalTempFilename(&fname));
196   TF_ASSERT_OK(
197       WriteStringToFile(env, fname, "line one\nline two\nline three\n"));
198   std::unique_ptr<RandomAccessFile> file;
199   TF_ASSERT_OK(env->NewRandomAccessFile(fname, &file));
200 
201   for (auto buf_size : BufferSizes()) {
202     std::unique_ptr<RandomAccessInputStream> input_stream(
203         new RandomAccessInputStream(file.get()));
204     BufferedInputStream in(input_stream.get(), buf_size);
205     string line;
206     TF_ASSERT_OK(in.SkipLine());
207     TF_ASSERT_OK(in.ReadLine(&line));
208     EXPECT_EQ(line, "line two");
209     TF_ASSERT_OK(in.SkipLine());
210     EXPECT_TRUE(errors::IsOutOfRange(in.SkipLine()));
211     // A second call should also return end of file
212     EXPECT_TRUE(errors::IsOutOfRange(in.SkipLine()));
213   }
214 }
215 
TEST(BufferedInputStream,SkipLine_NoTrailingNewLine)216 TEST(BufferedInputStream, SkipLine_NoTrailingNewLine) {
217   Env* env = Env::Default();
218   string fname;
219   ASSERT_TRUE(env->LocalTempFilename(&fname));
220   TF_ASSERT_OK(WriteStringToFile(env, fname, "line one\nline two\nline three"));
221   std::unique_ptr<RandomAccessFile> file;
222   TF_ASSERT_OK(env->NewRandomAccessFile(fname, &file));
223 
224   for (auto buf_size : BufferSizes()) {
225     std::unique_ptr<RandomAccessInputStream> input_stream(
226         new RandomAccessInputStream(file.get()));
227     BufferedInputStream in(input_stream.get(), buf_size);
228     string line;
229     TF_ASSERT_OK(in.SkipLine());
230     TF_ASSERT_OK(in.ReadLine(&line));
231     EXPECT_EQ(line, "line two");
232     TF_ASSERT_OK(in.SkipLine());
233     EXPECT_TRUE(errors::IsOutOfRange(in.SkipLine()));
234     // A second call should also return end of file
235     EXPECT_TRUE(errors::IsOutOfRange(in.SkipLine()));
236   }
237 }
238 
TEST(BufferedInputStream,SkipLine_EmptyLines)239 TEST(BufferedInputStream, SkipLine_EmptyLines) {
240   Env* env = Env::Default();
241   string fname;
242   ASSERT_TRUE(env->LocalTempFilename(&fname));
243   TF_ASSERT_OK(WriteStringToFile(env, fname, "line one\n\n\nline two"));
244   std::unique_ptr<RandomAccessFile> file;
245   TF_ASSERT_OK(env->NewRandomAccessFile(fname, &file));
246 
247   for (auto buf_size : BufferSizes()) {
248     std::unique_ptr<RandomAccessInputStream> input_stream(
249         new RandomAccessInputStream(file.get()));
250     BufferedInputStream in(input_stream.get(), buf_size);
251     string line;
252     TF_ASSERT_OK(in.SkipLine());
253     TF_ASSERT_OK(in.ReadLine(&line));
254     EXPECT_EQ(line, "");
255     TF_ASSERT_OK(in.SkipLine());
256     TF_ASSERT_OK(in.ReadLine(&line));
257     EXPECT_EQ(line, "line two");
258   }
259 }
260 
TEST(BufferedInputStream,ReadNBytes)261 TEST(BufferedInputStream, ReadNBytes) {
262   Env* env = Env::Default();
263   string fname;
264   ASSERT_TRUE(env->LocalTempFilename(&fname));
265   TF_ASSERT_OK(WriteStringToFile(env, fname, "0123456789"));
266   std::unique_ptr<RandomAccessFile> file;
267   TF_ASSERT_OK(env->NewRandomAccessFile(fname, &file));
268 
269   for (auto buf_size : BufferSizes()) {
270     std::unique_ptr<RandomAccessInputStream> input_stream(
271         new RandomAccessInputStream(file.get()));
272     tstring read;
273     BufferedInputStream in(input_stream.get(), buf_size);
274     EXPECT_EQ(0, in.Tell());
275     TF_ASSERT_OK(in.ReadNBytes(3, &read));
276     EXPECT_EQ(read, "012");
277     EXPECT_EQ(3, in.Tell());
278     TF_ASSERT_OK(in.ReadNBytes(0, &read));
279     EXPECT_EQ(read, "");
280     EXPECT_EQ(3, in.Tell());
281     TF_ASSERT_OK(in.ReadNBytes(4, &read));
282     EXPECT_EQ(read, "3456");
283     EXPECT_EQ(7, in.Tell());
284     TF_ASSERT_OK(in.ReadNBytes(0, &read));
285     EXPECT_EQ(read, "");
286     EXPECT_EQ(7, in.Tell());
287     EXPECT_TRUE(errors::IsOutOfRange(in.ReadNBytes(5, &read)));
288     EXPECT_EQ(read, "789");
289     EXPECT_EQ(10, in.Tell());
290     EXPECT_TRUE(errors::IsOutOfRange(in.ReadNBytes(5, &read)));
291     EXPECT_EQ(read, "");
292     EXPECT_EQ(10, in.Tell());
293     TF_ASSERT_OK(in.ReadNBytes(0, &read));
294     EXPECT_EQ(read, "");
295     EXPECT_EQ(10, in.Tell());
296   }
297 }
298 
TEST(BufferedInputStream,OutOfRangeCache)299 TEST(BufferedInputStream, OutOfRangeCache) {
300   for (auto buf_size : BufferSizes()) {
301     if (buf_size < 11) {
302       continue;
303     }
304     ReadOnceInputStream input_stream;
305     tstring read;
306     BufferedInputStream in(&input_stream, buf_size);
307     EXPECT_EQ(0, in.Tell());
308     TF_ASSERT_OK(in.ReadNBytes(3, &read));
309     EXPECT_EQ(read, "012");
310     EXPECT_EQ(3, in.Tell());
311     TF_ASSERT_OK((in.ReadNBytes(7, &read)));
312     EXPECT_EQ(read, "3456789");
313     EXPECT_EQ(10, in.Tell());
314     Status s = in.ReadNBytes(5, &read);
315     // Make sure the read is failing with OUT_OF_RANGE error. If it is failing
316     // with other errors, it is not caching the OUT_OF_RANGE properly.
317     EXPECT_EQ(error::OUT_OF_RANGE, s.code()) << s;
318     EXPECT_EQ(read, "");
319     // Empty read shouldn't cause an error even at the end of the file.
320     TF_ASSERT_OK(in.ReadNBytes(0, &read));
321     EXPECT_EQ(read, "");
322   }
323 }
324 
TEST(BufferedInputStream,SkipNBytes)325 TEST(BufferedInputStream, SkipNBytes) {
326   Env* env = Env::Default();
327   string fname;
328   ASSERT_TRUE(env->LocalTempFilename(&fname));
329   TF_ASSERT_OK(WriteStringToFile(env, fname, "0123456789"));
330   std::unique_ptr<RandomAccessFile> file;
331   TF_ASSERT_OK(env->NewRandomAccessFile(fname, &file));
332 
333   for (auto buf_size : BufferSizes()) {
334     std::unique_ptr<RandomAccessInputStream> input_stream(
335         new RandomAccessInputStream(file.get()));
336     tstring read;
337     BufferedInputStream in(input_stream.get(), buf_size);
338     EXPECT_EQ(0, in.Tell());
339     TF_ASSERT_OK(in.SkipNBytes(3));
340     EXPECT_EQ(3, in.Tell());
341     TF_ASSERT_OK(in.SkipNBytes(0));
342     EXPECT_EQ(3, in.Tell());
343     TF_ASSERT_OK(in.ReadNBytes(2, &read));
344     EXPECT_EQ(read, "34");
345     EXPECT_EQ(5, in.Tell());
346     TF_ASSERT_OK(in.SkipNBytes(0));
347     EXPECT_EQ(5, in.Tell());
348     TF_ASSERT_OK(in.SkipNBytes(2));
349     EXPECT_EQ(7, in.Tell());
350     TF_ASSERT_OK(in.ReadNBytes(1, &read));
351     EXPECT_EQ(read, "7");
352     EXPECT_EQ(8, in.Tell());
353     EXPECT_TRUE(errors::IsOutOfRange(in.SkipNBytes(5)));
354     EXPECT_EQ(10, in.Tell());
355     EXPECT_TRUE(errors::IsOutOfRange(in.SkipNBytes(5)));
356     EXPECT_EQ(10, in.Tell());
357     EXPECT_TRUE(errors::IsOutOfRange(in.ReadNBytes(5, &read)));
358     EXPECT_EQ(read, "");
359     EXPECT_EQ(10, in.Tell());
360   }
361 }
362 
TEST(BufferedInputStream,ReadNBytesRandomAccessFile)363 TEST(BufferedInputStream, ReadNBytesRandomAccessFile) {
364   Env* env = Env::Default();
365   string fname;
366   ASSERT_TRUE(env->LocalTempFilename(&fname));
367   TF_ASSERT_OK(WriteStringToFile(env, fname, "0123456789"));
368   std::unique_ptr<RandomAccessFile> file;
369   TF_ASSERT_OK(env->NewRandomAccessFile(fname, &file));
370 
371   for (auto buf_size : BufferSizes()) {
372     tstring read;
373     BufferedInputStream in(file.get(), buf_size);
374     EXPECT_EQ(0, in.Tell());
375     TF_ASSERT_OK(in.ReadNBytes(3, &read));
376     EXPECT_EQ(read, "012");
377     EXPECT_EQ(3, in.Tell());
378     TF_ASSERT_OK(in.ReadNBytes(0, &read));
379     EXPECT_EQ(read, "");
380     EXPECT_EQ(3, in.Tell());
381     TF_ASSERT_OK(in.ReadNBytes(4, &read));
382     EXPECT_EQ(read, "3456");
383     EXPECT_EQ(7, in.Tell());
384     TF_ASSERT_OK(in.ReadNBytes(0, &read));
385     EXPECT_EQ(read, "");
386     EXPECT_EQ(7, in.Tell());
387     EXPECT_TRUE(errors::IsOutOfRange(in.ReadNBytes(5, &read)));
388     EXPECT_EQ(read, "789");
389     EXPECT_EQ(10, in.Tell());
390     EXPECT_TRUE(errors::IsOutOfRange(in.ReadNBytes(5, &read)));
391     EXPECT_EQ(read, "");
392     EXPECT_EQ(10, in.Tell());
393     TF_ASSERT_OK(in.ReadNBytes(0, &read));
394     EXPECT_EQ(read, "");
395     EXPECT_EQ(10, in.Tell());
396   }
397 }
398 
TEST(BufferedInputStream,SkipNBytesRandomAccessFile)399 TEST(BufferedInputStream, SkipNBytesRandomAccessFile) {
400   Env* env = Env::Default();
401   string fname;
402   ASSERT_TRUE(env->LocalTempFilename(&fname));
403   TF_ASSERT_OK(WriteStringToFile(env, fname, "0123456789"));
404   std::unique_ptr<RandomAccessFile> file;
405   TF_ASSERT_OK(env->NewRandomAccessFile(fname, &file));
406 
407   for (auto buf_size : BufferSizes()) {
408     tstring read;
409     BufferedInputStream in(file.get(), buf_size);
410     EXPECT_EQ(0, in.Tell());
411     TF_ASSERT_OK(in.SkipNBytes(3));
412     EXPECT_EQ(3, in.Tell());
413     TF_ASSERT_OK(in.SkipNBytes(0));
414     EXPECT_EQ(3, in.Tell());
415     TF_ASSERT_OK(in.ReadNBytes(2, &read));
416     EXPECT_EQ(read, "34");
417     EXPECT_EQ(5, in.Tell());
418     TF_ASSERT_OK(in.SkipNBytes(0));
419     EXPECT_EQ(5, in.Tell());
420     TF_ASSERT_OK(in.SkipNBytes(2));
421     EXPECT_EQ(7, in.Tell());
422     TF_ASSERT_OK(in.ReadNBytes(1, &read));
423     EXPECT_EQ(read, "7");
424     EXPECT_EQ(8, in.Tell());
425     EXPECT_TRUE(errors::IsOutOfRange(in.SkipNBytes(5)));
426     EXPECT_EQ(10, in.Tell());
427     EXPECT_TRUE(errors::IsOutOfRange(in.SkipNBytes(5)));
428     EXPECT_EQ(10, in.Tell());
429     EXPECT_TRUE(errors::IsOutOfRange(in.ReadNBytes(5, &read)));
430     EXPECT_EQ(read, "");
431     EXPECT_EQ(10, in.Tell());
432   }
433 }
434 
TEST(BufferedInputStream,Seek)435 TEST(BufferedInputStream, Seek) {
436   Env* env = Env::Default();
437   string fname;
438   ASSERT_TRUE(env->LocalTempFilename(&fname));
439   TF_ASSERT_OK(WriteStringToFile(env, fname, "0123456789"));
440   std::unique_ptr<RandomAccessFile> file;
441   TF_ASSERT_OK(env->NewRandomAccessFile(fname, &file));
442 
443   for (auto buf_size : BufferSizes()) {
444     std::unique_ptr<RandomAccessInputStream> input_stream(
445         new RandomAccessInputStream(file.get()));
446     tstring read;
447     BufferedInputStream in(input_stream.get(), buf_size);
448 
449     // Seek forward
450     TF_ASSERT_OK(in.Seek(3));
451     EXPECT_EQ(3, in.Tell());
452 
453     // Read 4 bytes
454     TF_ASSERT_OK(in.ReadNBytes(4, &read));
455     EXPECT_EQ(read, "3456");
456     EXPECT_EQ(7, in.Tell());
457 
458     // Seek backwards
459     TF_ASSERT_OK(in.Seek(1));
460     TF_ASSERT_OK(in.ReadNBytes(4, &read));
461     EXPECT_EQ(read, "1234");
462     EXPECT_EQ(5, in.Tell());
463   }
464 }
465 
TEST(BufferedInputStream,Seek_NotReset)466 TEST(BufferedInputStream, Seek_NotReset) {
467   // This test verifies seek backwards within the buffer doesn't reset
468   // input_stream
469   Env* env = Env::Default();
470   string fname;
471   ASSERT_TRUE(env->LocalTempFilename(&fname));
472   TF_ASSERT_OK(WriteStringToFile(env, fname, "0123456789"));
473   std::unique_ptr<RandomAccessFile> file;
474   TF_ASSERT_OK(env->NewRandomAccessFile(fname, &file));
475 
476   std::unique_ptr<RandomAccessInputStream> input_stream(
477       new RandomAccessInputStream(file.get()));
478   tstring read;
479   BufferedInputStream in(input_stream.get(), 3);
480 
481   TF_ASSERT_OK(in.ReadNBytes(4, &read));
482   int before_tell = input_stream.get()->Tell();
483   EXPECT_EQ(before_tell, 6);
484   // Seek backwards
485   TF_ASSERT_OK(in.Seek(3));
486   int after_tell = input_stream.get()->Tell();
487   EXPECT_EQ(before_tell, after_tell);
488 }
489 
TEST(BufferedInputStream,ReadAll_Empty)490 TEST(BufferedInputStream, ReadAll_Empty) {
491   Env* env = Env::Default();
492   string fname;
493   ASSERT_TRUE(env->LocalTempFilename(&fname));
494   const string expected = "";
495   TF_ASSERT_OK(WriteStringToFile(env, fname, expected));
496   std::unique_ptr<RandomAccessFile> file;
497   TF_ASSERT_OK(env->NewRandomAccessFile(fname, &file));
498 
499   for (auto buf_size : BufferSizes()) {
500     RandomAccessInputStream input_stream(file.get());
501     BufferedInputStream in(&input_stream, buf_size);
502     string contents;
503     TF_ASSERT_OK(in.ReadAll(&contents));
504     EXPECT_EQ(expected, contents);
505   }
506 }
507 
TEST(BufferedInputStream,ReadAll_Text)508 TEST(BufferedInputStream, ReadAll_Text) {
509   Env* env = Env::Default();
510   string fname;
511   ASSERT_TRUE(env->LocalTempFilename(&fname));
512   const string expected = "line one\nline two\nline three";
513   TF_ASSERT_OK(WriteStringToFile(env, fname, expected));
514   std::unique_ptr<RandomAccessFile> file;
515   TF_ASSERT_OK(env->NewRandomAccessFile(fname, &file));
516 
517   for (auto buf_size : BufferSizes()) {
518     RandomAccessInputStream input_stream(file.get());
519     BufferedInputStream in(&input_stream, buf_size);
520     string contents;
521     TF_ASSERT_OK(in.ReadAll(&contents));
522     EXPECT_EQ(expected, contents);
523   }
524 }
525 
BM_BufferedReaderSmallReads(::testing::benchmark::State & state)526 void BM_BufferedReaderSmallReads(::testing::benchmark::State& state) {
527   const int buff_size = state.range(0);
528   const int file_size = state.range(1);
529   Env* env = Env::Default();
530   string fname;
531   ASSERT_TRUE(env->LocalTempFilename(&fname));
532 
533   const string file_elem = "0123456789";
534   std::unique_ptr<WritableFile> write_file;
535   TF_ASSERT_OK(env->NewWritableFile(fname, &write_file));
536   for (int i = 0; i < file_size; ++i) {
537     TF_ASSERT_OK(write_file->Append(file_elem));
538   }
539   TF_ASSERT_OK(write_file->Close());
540 
541   std::unique_ptr<RandomAccessFile> file;
542   TF_ASSERT_OK(env->NewRandomAccessFile(fname, &file));
543 
544   tstring result;
545 
546   int itr = 0;
547   for (auto s : state) {
548     BufferedInputStream in(file.get(), buff_size);
549     for (int64_t i = 0; i < 10 * file_size; ++i) {
550       TF_ASSERT_OK(in.ReadNBytes(1, &result))
551           << "i: " << i << " itr: " << itr << " buff_size: " << buff_size
552           << " file size: " << file_size;
553     }
554     ++itr;
555   }
556 }
557 BENCHMARK(BM_BufferedReaderSmallReads)
558     ->ArgPair(1, 5)
559     ->ArgPair(1, 1024)
560     ->ArgPair(10, 5)
561     ->ArgPair(10, 1024)
562     ->ArgPair(1024, 1024)
563     ->ArgPair(1024 * 1024, 1024)
564     ->ArgPair(1024 * 1024, 1024 * 1024)
565     ->ArgPair(256 * 1024 * 1024, 1024);
566 
567 }  // anonymous namespace
568 }  // namespace io
569 }  // namespace tensorflow
570