subprocess_harness.hpp

#ifndef BOLT_SUBPROCESS_HARNESS_HPP
#define BOLT_SUBPROCESS_HARNESS_HPP
#include <gtest/gtest.h>
#include <glog/logging.h>
#include <folly/Subprocess.h>
#include <folly/String.h>
#include <fstream>
#include <vector>
#include <string>
#include <algorithm>
#include <memory>
#include <thread>
#include <unistd.h>
#include <boost/filesystem.hpp>
#include "bolt/utils/osutils.hpp"

class SubprocessHarness : public ::testing::Test {
 protected:
  SubprocessHarness(const std::string &envFile /* for env vars   */,
                    const std::string &cmd     /* command to run */,
                    const std::string &name    /* for tmp dir    */,
                    const int sleepDuration    /* sleep duration */):
    tmpDir_(bolt::makeTempDir(boost::filesystem::temp_directory_path().string() +
                              "/" + name + ".XXXXXX")),
    command_(cmd),
    sleepDuration_(sleepDuration) {

    initializeEnv(envFile);
    env_.emplace(env_.begin(), "TMP_DIR=\"" + tmpDir_ + "\"");
  }

  ~SubprocessHarness() {
    if (subprocess_->returnCode().exited() ||
        subprocess_->returnCode().killed()) {
      return;
    }
    subprocess_->sendSignal(9);
    subprocess_->wait();
  }

  virtual void SetUp() {
    subprocess_ = std::make_shared<folly::Subprocess>(command_,
      folly::Subprocess::Options(), &env_);
    sleep(sleepDuration_);
  }

  virtual void TearDown() {
    subprocess_->sendSignal(9);
    subprocess_->wait();
  }

  void initializeEnv(std::string envFile) {
    std::ifstream file(envFile);

    if(file) {
      std::string envContents;
      file.seekg(0, file.end);
      envContents.reserve(file.tellg());
      file.seekg(0, file.beg);
      envContents.assign((std::istreambuf_iterator<char>(file)),
                         (std::istreambuf_iterator<char>()));

      folly::split('\n', envContents, env_);
      std::remove_if(env_.begin(),
                     env_.end(),
                     [](std::string & str) {
                       return str.size() == 0 || str[0] == '#';
                     });
    } else {
      throw std::runtime_error("Failed to open " + envFile);
    }
  }

  std::string tmpDir_;
  std::string command_;
  std::vector<std::string> env_;
  std::shared_ptr<folly::Subprocess> subprocess_;
  int sleepDuration_;
};
#endif