3 files changed, 108 insertions, 71 deletions

diff --git a/src/disk_interface_test.cc b/src/disk_interface_test.cc
index c431e75..b8ad5e6 100644
--- a/src/disk_interface_test.cc
+++ b/src/disk_interface_test.cc
@@ -23,87 +23,20 @@
 #include "graph.h"
 #include "test.h"
 
-using namespace std;
-
 namespace {
 
-#ifdef _WIN32
-#ifndef _mktemp_s
-/// mingw has no mktemp.  Implement one with the same type as the one
-/// found in the Windows API.
-int _mktemp_s(char* templ) {
-  char* ofs = strchr(templ, 'X');
-  sprintf(ofs, "%d", rand() % 1000000);
-  return 0;
-}
-#endif
-
-/// Windows has no mkdtemp.  Implement it in terms of _mktemp_s.
-char* mkdtemp(char* name_template) {
-  int err = _mktemp_s(name_template);
-  if (err < 0) {
-    perror("_mktemp_s");
-    return NULL;
-  }
-
-  err = _mkdir(name_template);
-  if (err < 0) {
-    perror("mkdir");
-    return NULL;
-  }
-
-  return name_template;
-}
-#endif
-
 class DiskInterfaceTest : public testing::Test {
  public:
   virtual void SetUp() {
-    // Because we do real disk accesses, we create a temp dir within
-    // the system temporary directory.
-
-    // First change into the system temp dir and save it for cleanup.
-    start_dir_ = GetSystemTempDir();
-    ASSERT_EQ(0, chdir(start_dir_.c_str()));
-
-    // Then create and change into a temporary subdirectory of that.
-    temp_dir_name_ = MakeTempDir();
-    ASSERT_FALSE(temp_dir_name_.empty());
-    ASSERT_EQ(0, chdir(temp_dir_name_.c_str()));
+    // These tests do real disk accesses, so create a temp dir.
+    temp_dir_.CreateAndEnter("Ninja-DiskInterfaceTest");
   }
 
   virtual void TearDown() {
-    // Move out of the directory we're about to clobber.
-    ASSERT_EQ(0, chdir(start_dir_.c_str()));
-#ifdef _WIN32
-    ASSERT_EQ(0, system(("rmdir /s /q " + temp_dir_name_).c_str()));
-#else
-    ASSERT_EQ(0, system(("rm -rf " + temp_dir_name_).c_str()));
-#endif
-  }
-
-  string GetSystemTempDir() {
-#ifdef _WIN32
-    char buf[1024];
-    if (!GetTempPath(sizeof(buf), buf))
-      return "";
-    return buf;
-#else
-    const char* tempdir = getenv("TMPDIR");
-    if (tempdir)
-      return tempdir;
-    return "/tmp";
-#endif
-  }
-
-  string MakeTempDir() {
-    char name_template[] = "DiskInterfaceTest-XXXXXX";
-    char* name = mkdtemp(name_template);
-    return name ? name : "";
+    temp_dir_.Cleanup();
   }
 
-  string start_dir_;
-  string temp_dir_name_;
+  ScopedTempDir temp_dir_;
   RealDiskInterface disk_;
 };
 
diff --git a/src/test.cc b/src/test.cc
index 925a6cb..f5b037f 100644
--- a/src/test.cc
+++ b/src/test.cc
@@ -16,7 +16,57 @@
 
 #include <algorithm>
 
+#include <errno.h>
+
 #include "parsers.h"
+#include "util.h"
+
+namespace {
+
+#ifdef _WIN32
+#ifndef _mktemp_s
+/// mingw has no mktemp.  Implement one with the same type as the one
+/// found in the Windows API.
+int _mktemp_s(char* templ) {
+  char* ofs = strchr(templ, 'X');
+  sprintf(ofs, "%d", rand() % 1000000);
+  return 0;
+}
+#endif
+
+/// Windows has no mkdtemp.  Implement it in terms of _mktemp_s.
+char* mkdtemp(char* name_template) {
+  int err = _mktemp_s(name_template);
+  if (err < 0) {
+    perror("_mktemp_s");
+    return NULL;
+  }
+
+  err = _mkdir(name_template);
+  if (err < 0) {
+    perror("mkdir");
+    return NULL;
+  }
+
+  return name_template;
+}
+#endif  // WIN32
+
+string GetSystemTempDir() {
+#ifdef _WIN32
+  char buf[1024];
+  if (!GetTempPath(sizeof(buf), buf))
+    return "";
+  return buf;
+#else
+  const char* tempdir = getenv("TMPDIR");
+  if (tempdir)
+    return tempdir;
+  return "/tmp";
+#endif
+}
+
+}  // anonymous namespace
 
 StateTestWithBuiltinRules::StateTestWithBuiltinRules() {
   AssertParse(&state_,
@@ -74,3 +124,44 @@ int VirtualFileSystem::RemoveFile(const string& path) {
     return 1;
   }
 }
+
+void ScopedTempDir::CreateAndEnter(const string& name) {
+  // First change into the system temp dir and save it for cleanup.
+  start_dir_ = GetSystemTempDir();
+  if (start_dir_.empty())
+    Fatal("couldn't get system temp dir");
+  if (chdir(start_dir_.c_str()) < 0)
+    Fatal("chdir: %s", strerror(errno));
+
+  // Create a temporary subdirectory of that.
+  char name_template[1024];
+  strcpy(name_template, name.c_str());
+  strcat(name_template, "-XXXXXX");
+  char* tempname = mkdtemp(name_template);
+  if (!tempname)
+    Fatal("mkdtemp: %s", strerror(errno));
+  temp_dir_name_ = tempname;
+
+  // chdir into the new temporary directory.
+  if (chdir(temp_dir_name_.c_str()) < 0)
+    Fatal("chdir: %s", strerror(errno));
+}
+
+void ScopedTempDir::Cleanup() {
+  if (temp_dir_name_.empty())
+    return;  // Something went wrong earlier.
+
+  // Move out of the directory we're about to clobber.
+  if (chdir(start_dir_.c_str()) < 0)
+    Fatal("chdir: %s", strerror(errno));
+
+#ifdef _WIN32
+  string command = "rmdir /s /q " + temp_dir_name_;
+#else
+  string command = "rm -rf " + temp_dir_name_;
+#endif
+  if (system(command.c_str()) < 0)
+    Fatal("system: %s", strerror(errno));
+
+  temp_dir_name_.clear();
+}
diff --git a/src/test.h b/src/test.h
index c8c2758..df544c6 100644
--- a/src/test.h
+++ b/src/test.h
@@ -61,4 +61,17 @@ struct VirtualFileSystem : public DiskInterface {
   set<string> files_removed_;
 };
 
+struct ScopedTempDir {
+  /// Create a temporary directory and chdir into it.
+  void CreateAndEnter(const string& name);
+
+  /// Clean up the temporary directory.
+  void Cleanup();
+
+  /// The temp directory containing our dir.
+  string start_dir_;
+  /// The subdirectory name for our dir, or empty if it hasn't been set up.
+  string temp_dir_name_;
+};
+
 #endif // NINJA_TEST_H_