From 1b08746e197bd5231fc5151d77385fcd2d7afd79 Mon Sep 17 00:00:00 2001
From: Yedidya Feldblum <yfeldblum@fb.com>
Date: Tue, 18 Aug 2015 02:09:11 -0700
Subject: [PATCH] EnvVarSaver.

Summary: [Folly] EnvVarSaver.

    TEST(SomeBigClassTest, ChangesEnvVars) {
      folly::EnvVarSaver saver;
      setenv("USER", "root", 1);
      BigClass().doSomethingWithUser();
    }

Reviewed By: @Gownta

Differential Revision: D2354679
---
 folly/experimental/TestUtil.cpp          | 31 +++++++++++++++++++
 folly/experimental/TestUtil.h            |  9 ++++++
 folly/experimental/test/TestUtilTest.cpp | 38 ++++++++++++++++++++++++
 3 files changed, 78 insertions(+)

diff --git a/folly/experimental/TestUtil.cpp b/folly/experimental/TestUtil.cpp
index ae84cc31..4ff1cf32 100644
--- a/folly/experimental/TestUtil.cpp
+++ b/folly/experimental/TestUtil.cpp
@@ -26,6 +26,7 @@
 #include <folly/Exception.h>
 #include <folly/File.h>
 #include <folly/FileUtil.h>
+#include <folly/String.h>
 
 namespace folly {
 namespace test {
@@ -179,5 +180,35 @@ std::string CaptureFD::readIncremental() {
   return std::string(buf.get(), size);
 }
 
+static std::map<std::string, std::string> getEnvVarMap() {
+  std::map<std::string, std::string> data;
+  for (auto it = environ; *it != nullptr; ++it) {
+    std::string key, value;
+    split("=", *it, key, value);
+    if (key.empty()) {
+      continue;
+    }
+    CHECK(!data.count(key)) << "already contains: " << key;
+    data.emplace(move(key), move(value));
+  }
+  return data;
+}
+
+EnvVarSaver::EnvVarSaver() {
+  saved_ = getEnvVarMap();
+}
+
+EnvVarSaver::~EnvVarSaver() {
+  for (const auto& kvp : getEnvVarMap()) {
+    if (saved_.count(kvp.first)) {
+      continue;
+    }
+    PCHECK(0 == unsetenv(kvp.first.c_str()));
+  }
+  for (const auto& kvp : saved_) {
+    PCHECK(0 == setenv(kvp.first.c_str(), kvp.second.c_str(), (int)true));
+  }
+}
+
 }  // namespace test
 }  // namespace folly
diff --git a/folly/experimental/TestUtil.h b/folly/experimental/TestUtil.h
index dbf84768..e1bd1b9b 100644
--- a/folly/experimental/TestUtil.h
+++ b/folly/experimental/TestUtil.h
@@ -17,6 +17,7 @@
 #ifndef FOLLY_TESTUTIL_H_
 #define FOLLY_TESTUTIL_H_
 
+#include <map>
 #include <string>
 #include <folly/Range.h>
 #include <folly/experimental/io/FsUtil.h>
@@ -191,6 +192,14 @@ private:
   off_t readOffset_;  // for incremental reading
 };
 
+class EnvVarSaver {
+public:
+  EnvVarSaver();
+  ~EnvVarSaver();
+private:
+  std::map<std::string, std::string> saved_;
+};
+
 }  // namespace test
 }  // namespace folly
 
diff --git a/folly/experimental/test/TestUtilTest.cpp b/folly/experimental/test/TestUtilTest.cpp
index d434c9b7..af42a0ca 100644
--- a/folly/experimental/test/TestUtilTest.cpp
+++ b/folly/experimental/test/TestUtilTest.cpp
@@ -23,6 +23,7 @@
 #include <system_error>
 
 #include <boost/algorithm/string.hpp>
+#include <folly/Memory.h>
 #include <glog/logging.h>
 #include <gtest/gtest.h>
 
@@ -135,6 +136,43 @@ TEST(CaptureFD, GlogPatterns) {
   }
 }
 
+class EnvVarSaverTest : public testing::Test {};
+
+TEST_F(EnvVarSaverTest, ExampleNew) {
+  auto key = "hahahahaha";
+  EXPECT_EQ(nullptr, getenv(key));
+
+  auto saver = make_unique<EnvVarSaver>();
+  PCHECK(0 == setenv(key, "blah", true));
+  EXPECT_EQ("blah", std::string{getenv(key)});
+  saver = nullptr;
+  EXPECT_EQ(nullptr, getenv(key));
+}
+
+TEST_F(EnvVarSaverTest, ExampleExisting) {
+  auto key = "USER";
+  EXPECT_NE(nullptr, getenv(key));
+  auto value = std::string{getenv(key)};
+
+  auto saver = make_unique<EnvVarSaver>();
+  PCHECK(0 == setenv(key, "blah", true));
+  EXPECT_EQ("blah", std::string{getenv(key)});
+  saver = nullptr;
+  EXPECT_TRUE(value == getenv(key));
+}
+
+TEST_F(EnvVarSaverTest, ExampleDeleting) {
+  auto key = "USER";
+  EXPECT_NE(nullptr, getenv(key));
+  auto value = std::string{getenv(key)};
+
+  auto saver = make_unique<EnvVarSaver>();
+  PCHECK(0 == unsetenv(key));
+  EXPECT_EQ(nullptr, getenv(key));
+  saver = nullptr;
+  EXPECT_TRUE(value == getenv(key));
+}
+
 int main(int argc, char *argv[]) {
   testing::InitGoogleTest(&argc, argv);
   gflags::ParseCommandLineFlags(&argc, &argv, true);
-- 
2.34.1