From 338a260317f0e7926df1fb01d8dc432c8b07cb57 Mon Sep 17 00:00:00 2001
From: Alexey Spiridonov <lesha@fb.com>
Date: Fri, 2 Oct 2015 22:30:41 -0700
Subject: [PATCH] Add "consume all captured output" callback to CaptureFD

Summary: I noticed myself trying to fake this kind of callback for a log-based test I was writing. It seems much nicer to add the callback to `CaptureFD` than roll ugly wrappers around it to do the same thing.

Reviewed By: @yfeldblum

Differential Revision: D2506106
---
 folly/experimental/TestUtil.cpp          |  7 ++++--
 folly/experimental/TestUtil.h            | 17 +++++++++++---
 folly/experimental/test/TestUtilTest.cpp | 28 ++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/folly/experimental/TestUtil.cpp b/folly/experimental/TestUtil.cpp
index dbdd7bd7..7d7ba81e 100644
--- a/folly/experimental/TestUtil.cpp
+++ b/folly/experimental/TestUtil.cpp
@@ -142,7 +142,8 @@ bool hasNoPCREPatternMatch(StringPiece pattern, StringPiece target) {
 
 }  // namespace detail
 
-CaptureFD::CaptureFD(int fd) : fd_(fd), readOffset_(0) {
+CaptureFD::CaptureFD(int fd, ChunkCob chunk_cob)
+    : chunkCob_(std::move(chunk_cob)), fd_(fd), readOffset_(0) {
   oldFDCopy_ = dup(fd_);
   PCHECK(oldFDCopy_ != -1) << "Could not copy FD " << fd_;
 
@@ -154,6 +155,7 @@ CaptureFD::CaptureFD(int fd) : fd_(fd), readOffset_(0) {
 
 void CaptureFD::release() {
   if (oldFDCopy_ != fd_) {
+    readIncremental();  // Feed chunkCob_
     PCHECK(dup2(oldFDCopy_, fd_) != -1) << "Could not restore old FD "
       << oldFDCopy_ << " into " << fd_;
     PCHECK(close(oldFDCopy_) != -1) << "Could not close " << oldFDCopy_;
@@ -165,7 +167,7 @@ CaptureFD::~CaptureFD() {
   release();
 }
 
-std::string CaptureFD::read() {
+std::string CaptureFD::read() const {
   std::string contents;
   std::string filename = file_.path().string();
   PCHECK(folly::readFile(filename.c_str(), contents));
@@ -181,6 +183,7 @@ std::string CaptureFD::readIncremental() {
   auto bytes_read = folly::preadFull(f.fd(), buf.get(), size, readOffset_);
   PCHECK(size == bytes_read);
   readOffset_ += size;
+  chunkCob_(StringPiece(buf.get(), buf.get() + size));
   return std::string(buf.get(), size);
 }
 
diff --git a/folly/experimental/TestUtil.h b/folly/experimental/TestUtil.h
index e1bd1b9b..3fa6a9f5 100644
--- a/folly/experimental/TestUtil.h
+++ b/folly/experimental/TestUtil.h
@@ -157,12 +157,22 @@ inline std::string glogErrOrWarnPattern() { return ".*(^|\n)[EW][0-9].*"; }
 
 /**
  * Temporarily capture a file descriptor by redirecting it into a file.
- * You can consume its output either all-at-once or incrementally.
+ * You can consume its entire output thus far via read(), incrementally
+ * via readIncremental(), or via callback using chunk_cob.
  * Great for testing logging (see also glog*Pattern()).
  */
 class CaptureFD {
+private:
+  struct NoOpChunkCob { void operator()(StringPiece) {} };
 public:
-  explicit CaptureFD(int fd);
+  using ChunkCob = std::function<void(folly::StringPiece)>;
+
+  /**
+   * chunk_cob is is guaranteed to consume all the captured output. It is
+   * invoked on each readIncremental(), and also on FD release to capture
+   * as-yet unread lines.  Chunks can be empty.
+   */
+  explicit CaptureFD(int fd, ChunkCob chunk_cob = NoOpChunkCob());
   ~CaptureFD();
 
   /**
@@ -175,7 +185,7 @@ public:
   /**
    * Reads the whole file into a string, but does not remove the redirect.
    */
-  std::string read();
+  std::string read() const;
 
   /**
    * Read any bytes that were appended to the file since the last
@@ -184,6 +194,7 @@ public:
   std::string readIncremental();
 
 private:
+  ChunkCob chunkCob_;
   TemporaryFile file_;
 
   int fd_;
diff --git a/folly/experimental/test/TestUtilTest.cpp b/folly/experimental/test/TestUtilTest.cpp
index af42a0ca..aca97e83 100644
--- a/folly/experimental/test/TestUtilTest.cpp
+++ b/folly/experimental/test/TestUtilTest.cpp
@@ -136,6 +136,34 @@ TEST(CaptureFD, GlogPatterns) {
   }
 }
 
+TEST(CaptureFD, ChunkCob) {
+  std::vector<std::string> chunks;
+  {
+    CaptureFD stderr(2, [&](StringPiece p) {
+      chunks.emplace_back(p.str());
+      switch (chunks.size()) {
+        case 1:
+          EXPECT_PCRE_MATCH(".*foo.*bar.*", p);
+          break;
+        case 2:
+          EXPECT_PCRE_MATCH("[^\n]*baz.*", p);
+          break;
+        default:
+          FAIL() << "Got too many chunks: " << chunks.size();
+      }
+    });
+    LOG(INFO) << "foo";
+    LOG(INFO) << "bar";
+    EXPECT_PCRE_MATCH(".*foo.*bar.*", stderr.read());
+    auto chunk = stderr.readIncremental();
+    EXPECT_EQ(chunks.at(0), chunk);
+    LOG(INFO) << "baz";
+    EXPECT_PCRE_MATCH(".*foo.*bar.*baz.*", stderr.read());
+  }
+  EXPECT_EQ(2, chunks.size());
+}
+
+
 class EnvVarSaverTest : public testing::Test {};
 
 TEST_F(EnvVarSaverTest, ExampleNew) {
-- 
2.34.1