2017

[folly.git] / folly / gen / test / StringTest.cpp

diff --git a/folly/gen/test/StringTest.cpp b/folly/gen/test/StringTest.cpp
index 71acab88418bdd696e96d5d1aa7eb15374056033..82239479a2a662c992821c727fe68e2f8cd93144 100644 (file)
--- a/folly/gen/test/StringTest.cpp
+++ b/folly/gen/test/StringTest.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright 2014 Facebook, Inc.
+ * Copyright 2017 Facebook, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -15,12 +15,13 @@
  */
 
 #include <glog/logging.h>
-#include <gtest/gtest.h>
+
 #include <iosfwd>
 #include <map>
 #include <vector>
 
 #include <folly/gen/String.h>
+#include <folly/portability/GTest.h>
 
 using namespace folly::gen;
 using namespace folly;
@@ -104,13 +105,15 @@ TEST(StringGen, Split) {
 TEST(StringGen, SplitByNewLine) {
   auto collect = eachTo<std::string>() | as<vector>();
   {
-    auto pieces = lines("hello\n\n world\r\n goodbye\r meow") | collect;
-    EXPECT_EQ(5, pieces.size());
+    auto pieces = lines("hello\n\n world\r\n goodbye\r me\n\row") | collect;
+    EXPECT_EQ(7, pieces.size());
     EXPECT_EQ("hello", pieces[0]);
     EXPECT_EQ("", pieces[1]);
     EXPECT_EQ(" world", pieces[2]);
     EXPECT_EQ(" goodbye", pieces[3]);
-    EXPECT_EQ(" meow", pieces[4]);
+    EXPECT_EQ(" me", pieces[4]);
+    EXPECT_EQ("", pieces[5]);
+    EXPECT_EQ("ow", pieces[6]);
   }
 }
 
@@ -258,6 +261,50 @@ TEST(StringGen, Resplit) {
   }
 }
 
+void checkResplitMaxLength(vector<string> ins,
+                           char delim,
+                           uint64_t maxLength,
+                           vector<string> outs) {
+  vector<std::string> pieces;
+  auto splitter = streamSplitter(delim, [&pieces](StringPiece s) {
+    pieces.push_back(string(s.begin(), s.end()));
+    return true;
+  }, maxLength);
+  for (const auto& in : ins) {
+    splitter(in);
+  }
+  splitter.flush();
+
+  EXPECT_EQ(outs.size(), pieces.size());
+  for (size_t i = 0; i < outs.size(); ++i) {
+    EXPECT_EQ(outs[i], pieces[i]);
+  }
+
+  // Also check the concatenated input against the same output
+  if (ins.size() > 1) {
+    checkResplitMaxLength({folly::join("", ins)}, delim, maxLength, outs);
+  }
+}
+
+TEST(StringGen, ResplitMaxLength) {
+  checkResplitMaxLength(
+    {"hel", "lo,", ", world", ", goodbye, m", "ew"}, ',', 5,
+    {"hello", ",", ",", " worl", "d,", " good", "bye,", " mew"}
+  );
+  // " meow" cannot be "end of stream", since it's maxLength long
+  checkResplitMaxLength(
+    {"hel", "lo,", ", world", ", goodbye, m", "eow"}, ',', 5,
+    {"hello", ",", ",", " worl", "d,", " good", "bye,", " meow", ""}
+  );
+  checkResplitMaxLength(
+    {"||", "", "", "", "|a|b", "cdefghijklmn", "|opqrst",
+     "uvwx|y|||", "z", "0123456789", "|", ""}, '|', 2,
+    {"|", "|", "|", "a|", "bc", "de", "fg", "hi", "jk", "lm", "n|", "op", "qr",
+     "st", "uv", "wx", "|", "y|", "|", "|", "z0", "12", "34", "56", "78", "9|",
+     ""}
+  );
+}
+
 template<typename F>
 void runUnsplitSuite(F fn) {
   fn("hello, world");
@@ -302,8 +349,16 @@ TEST(StringGen, Unsplit) {
   EXPECT_EQ("1, 2, 3", seq(1, 3) | unsplit(", "));
 }
 
-int main(int argc, char *argv[]) {
-  testing::InitGoogleTest(&argc, argv);
-  google::ParseCommandLineFlags(&argc, &argv, true);
-  return RUN_ALL_TESTS();
+TEST(StringGen, Batch) {
+  std::vector<std::string> chunks{
+      "on", "e\nt", "w", "o", "\nthr", "ee\nfo", "ur\n",
+  };
+  std::vector<std::string> lines{
+      "one", "two", "three", "four",
+  };
+  EXPECT_EQ(4, from(chunks) | resplit('\n') | count);
+  EXPECT_EQ(4, from(chunks) | resplit('\n') | batch(2) | rconcat | count);
+  EXPECT_EQ(4, from(chunks) | resplit('\n') | batch(3) | rconcat | count);
+  EXPECT_EQ(lines, from(chunks) | resplit('\n') | eachTo<std::string>() |
+                       batch(3) | rconcat | as<vector>());
 }