From 010dba3495c7e761fe24f77ec1af3b5e08d69ca1 Mon Sep 17 00:00:00 2001
From: Kane York <kanepyork@gmail.com>
Date: Fri, 14 Aug 2020 17:23:12 -0700
Subject: chore(tvix/tests): Add OpenTemporaryStore() helper

Change-Id: Ia655175c255f9cf5a47e4e7a20373a4ba4315ed9
Reviewed-on: https://cl.tvl.fyi/c/depot/+/1753
Tested-by: BuildkiteCI
Reviewed-by: glittershark <grfn@gws.fyi>
---
 third_party/nix/src/tests/store-util.hh  | 73 ++++++++++++++++++++++++++++++++
 third_party/nix/src/tests/store_tests.cc |  8 +---
 2 files changed, 74 insertions(+), 7 deletions(-)
 create mode 100644 third_party/nix/src/tests/store-util.hh

diff --git a/third_party/nix/src/tests/store-util.hh b/third_party/nix/src/tests/store-util.hh
new file mode 100644
index 0000000000..170794616b
--- /dev/null
+++ b/third_party/nix/src/tests/store-util.hh
@@ -0,0 +1,73 @@
+#pragma once
+
+#include <filesystem>
+
+#include <absl/status/statusor.h>
+#include <absl/strings/escaping.h>
+#include <glog/logging.h>
+#include <sys/random.h>
+
+#include "libstore/local-store.hh"
+
+namespace nix {
+
+class StoreTest : public ::testing::Test {
+ public:
+  virtual void TearDown() {
+    for (auto fn : cleanup_funcs_) {
+      try {
+        fn();
+      } catch (std::exception e) {
+        LOG(ERROR) << e.what();
+      }
+    }
+  }
+
+  absl::StatusOr<std::filesystem::path> OpenTempDir(
+      std::filesystem::path parent = std::filesystem::temp_directory_path()) {
+    for (;;) {
+      constexpr int kByteCnt = 9;
+      std::array<char, kByteCnt> randBytes;
+      if (getrandom(randBytes.data(), kByteCnt, 0) < 0) {
+        return absl::InternalError("getrandom() failed");
+      }
+      std::string suffix = absl::WebSafeBase64Escape(
+          absl::string_view(randBytes.data(), kByteCnt));
+      CHECK(suffix != "");
+
+      // Workaround for stdlib bug: use .assign() and ::errc
+      // https://stackoverflow.com/a/52401295/1210278
+      std::error_code ec_exists;
+      ec_exists.assign(EEXIST, std::system_category());
+
+      std::error_code ec;
+      std::filesystem::path candidate =
+          parent / absl::StrCat("nixtest-", suffix);
+      if (std::filesystem::create_directory(candidate, ec)) {
+        cleanup_funcs_.push_back(
+            [candidate]() { std::filesystem::remove_all(candidate); });
+        return candidate;
+      } else if (ec == ec_exists || ec == std::errc::file_exists) {
+        // Directory existed, retry
+        continue;
+      } else {
+        return absl::InternalError(absl::StrCat(
+            "could not create dir ", candidate.c_str(), ": ", ec.message()));
+      }
+    }
+  }
+
+  absl::StatusOr<std::unique_ptr<nix::LocalStore>> OpenTemporaryStore() {
+    ASSIGN_OR_RETURN(std::filesystem::path storePath, OpenTempDir());
+
+    nix::Store::Params params;
+    params["root"] = storePath;
+
+    return std::make_unique<nix::LocalStore>(params);
+  }
+
+ private:
+  std::vector<std::function<void(void)>> cleanup_funcs_;
+};
+
+}  // namespace nix
diff --git a/third_party/nix/src/tests/store_tests.cc b/third_party/nix/src/tests/store_tests.cc
index 1aea33bfd1..936a0cee1b 100644
--- a/third_party/nix/src/tests/store_tests.cc
+++ b/third_party/nix/src/tests/store_tests.cc
@@ -10,6 +10,7 @@
 
 #include "libstore/binary-cache-store.hh"
 #include "libstore/mock-binary-cache-store.hh"
+#include "tests/store-util.hh"
 
 using ::testing::HasSubstr;
 
@@ -17,13 +18,6 @@ namespace nix {
 
 MakeError(InjectedError, Error);
 
-class StoreTest : public ::testing::Test {
- public:
-  static void SetUpTestSuite() {
-    google::InitGoogleLogging("--logtostderr=false");
-  }
-};
-
 class BinaryCacheStoreTest : public StoreTest {};
 
 constexpr absl::string_view kXZHeader = "7zXZ";
-- 
cgit 1.4.1