KenLM 58da338b

1 files changed, 35 insertions, 20 deletions

diff --git a/klm/util/read_compressed_test.cc b/klm/util/read_compressed_test.cc
index 6fd97e5e..9cb4a4b9 100644
--- a/klm/util/read_compressed_test.cc
+++ b/klm/util/read_compressed_test.cc
@@ -25,19 +25,34 @@ void ReadLoop(ReadCompressed &reader, void *to_void, std::size_t amount) {
   }
 }
 
-void TestRandom(const char *compressor) {
-  const uint32_t kSize4 = 100000 / 4;
+const uint32_t kSize4 = 100000 / 4;
+
+std::string WriteRandom() {
   char name[] = "tempXXXXXX";
+  scoped_fd original(mkstemp(name));
+  BOOST_REQUIRE(original.get() > 0);
+  for (uint32_t i = 0; i < kSize4; ++i) {
+    WriteOrThrow(original.get(), &i, sizeof(uint32_t));
+  }
+  return name;
+}
 
-  // Write test file.  
-  {
-    scoped_fd original(mkstemp(name));
-    BOOST_REQUIRE(original.get() > 0);
-    for (uint32_t i = 0; i < kSize4; ++i) {
-      WriteOrThrow(original.get(), &i, sizeof(uint32_t));
-    }
+void VerifyRead(ReadCompressed &reader) {
+  for (uint32_t i = 0; i < kSize4; ++i) {
+    uint32_t got;
+    ReadLoop(reader, &got, sizeof(uint32_t));
+    BOOST_CHECK_EQUAL(i, got);
   }
 
+  char ignored;
+  BOOST_CHECK_EQUAL((std::size_t)0, reader.Read(&ignored, 1));
+  // Test double EOF call.
+  BOOST_CHECK_EQUAL((std::size_t)0, reader.Read(&ignored, 1));
+}
+
+void TestRandom(const char *compressor) {
+  std::string name(WriteRandom());
+
   char gzname[] = "tempXXXXXX";
   scoped_fd gzipped(mkstemp(gzname));
 
@@ -52,20 +67,11 @@ void TestRandom(const char *compressor) {
   command += "\"";
   BOOST_REQUIRE_EQUAL(0, system(command.c_str()));
 
-  BOOST_CHECK_EQUAL(0, unlink(name));
+  BOOST_CHECK_EQUAL(0, unlink(name.c_str()));
   BOOST_CHECK_EQUAL(0, unlink(gzname));
 
   ReadCompressed reader(gzipped.release());
-  for (uint32_t i = 0; i < kSize4; ++i) {
-    uint32_t got;
-    ReadLoop(reader, &got, sizeof(uint32_t));
-    BOOST_CHECK_EQUAL(i, got);
-  }
-
-  char ignored;
-  BOOST_CHECK_EQUAL((std::size_t)0, reader.Read(&ignored, 1));
-  // Test double EOF call.
-  BOOST_CHECK_EQUAL((std::size_t)0, reader.Read(&ignored, 1));
+  VerifyRead(reader);
 }
 
 BOOST_AUTO_TEST_CASE(Uncompressed) {
@@ -90,5 +96,14 @@ BOOST_AUTO_TEST_CASE(ReadXZ) {
 }
 #endif
 
+BOOST_AUTO_TEST_CASE(IStream) {
+  std::string name(WriteRandom());
+  std::fstream stream(name.c_str(), std::ios::in);
+  BOOST_CHECK_EQUAL(0, unlink(name.c_str()));
+  ReadCompressed reader;
+  reader.Reset(stream);
+  VerifyRead(reader);
+}
+
 } // namespace
 } // namespace util