1 files changed, 22 insertions, 24 deletions

diff --git a/tensorflow/core/platform/hadoop/hadoop_file_system_test.cc b/tensorflow/core/platform/hadoop/hadoop_file_system_test.cc
index cb12913f13..59e1d23645 100644
--- a/tensorflow/core/platform/hadoop/hadoop_file_system_test.cc
+++ b/tensorflow/core/platform/hadoop/hadoop_file_system_test.cc
@@ -28,6 +28,15 @@ class HadoopFileSystemTest : public ::testing::Test {
  protected:
   HadoopFileSystemTest() {}
 
+  string TmpDir(const string& path) {
+    char* test_dir = getenv("HADOOP_TEST_TMPDIR");
+    if (test_dir != nullptr) {
+      return io::JoinPath(string(test_dir), path);
+    } else {
+      return "file://" + io::JoinPath(testing::TmpDir(), path);
+    }
+  }
+
   Status WriteString(const string& fname, const string& content) {
     std::unique_ptr<WritableFile> writer;
     TF_RETURN_IF_ERROR(hdfs.NewWritableFile(fname, &writer));
@@ -58,8 +67,7 @@ class HadoopFileSystemTest : public ::testing::Test {
 };
 
 TEST_F(HadoopFileSystemTest, RandomAccessFile) {
-  const string fname =
-      "file://" + io::JoinPath(testing::TmpDir(), "RandomAccessFile");
+  const string fname = TmpDir("RandomAccessFile");
   const string content = "abcdefghijklmn";
   TF_ASSERT_OK(WriteString(fname, content));
 
@@ -83,8 +91,7 @@ TEST_F(HadoopFileSystemTest, RandomAccessFile) {
 
 TEST_F(HadoopFileSystemTest, WritableFile) {
   std::unique_ptr<WritableFile> writer;
-  const string fname =
-      "file://" + io::JoinPath(testing::TmpDir(), "WritableFile");
+  const string fname = TmpDir("WritableFile");
   TF_EXPECT_OK(hdfs.NewWritableFile(fname, &writer));
   TF_EXPECT_OK(writer->Append("content1,"));
   TF_EXPECT_OK(writer->Append("content2"));
@@ -98,16 +105,14 @@ TEST_F(HadoopFileSystemTest, WritableFile) {
 }
 
 TEST_F(HadoopFileSystemTest, FileExists) {
-  const string fname =
-      "file://" + io::JoinPath(testing::TmpDir(), "FileExists");
+  const string fname = TmpDir("FileExists");
   EXPECT_EQ(error::Code::NOT_FOUND, hdfs.FileExists(fname).code());
   TF_ASSERT_OK(WriteString(fname, "test"));
   TF_EXPECT_OK(hdfs.FileExists(fname));
 }
 
 TEST_F(HadoopFileSystemTest, GetChildren) {
-  const string base =
-      "file://" + io::JoinPath(testing::TmpDir(), "GetChildren");
+  const string base = TmpDir("GetChildren");
   TF_EXPECT_OK(hdfs.CreateDir(base));
 
   const string file = io::JoinPath(base, "testfile.csv");
@@ -122,16 +127,14 @@ TEST_F(HadoopFileSystemTest, GetChildren) {
 }
 
 TEST_F(HadoopFileSystemTest, DeleteFile) {
-  const string fname =
-      "file://" + io::JoinPath(testing::TmpDir(), "DeleteFile");
+  const string fname = TmpDir("DeleteFile");
   EXPECT_FALSE(hdfs.DeleteFile(fname).ok());
   TF_ASSERT_OK(WriteString(fname, "test"));
   TF_EXPECT_OK(hdfs.DeleteFile(fname));
 }
 
 TEST_F(HadoopFileSystemTest, GetFileSize) {
-  const string fname =
-      "file://" + io::JoinPath(testing::TmpDir(), "GetFileSize");
+  const string fname = TmpDir("GetFileSize");
   TF_ASSERT_OK(WriteString(fname, "test"));
   uint64 file_size = 0;
   TF_EXPECT_OK(hdfs.GetFileSize(fname, &file_size));
@@ -139,8 +142,7 @@ TEST_F(HadoopFileSystemTest, GetFileSize) {
 }
 
 TEST_F(HadoopFileSystemTest, CreateDirStat) {
-  const string dir =
-      "file://" + io::JoinPath(testing::TmpDir(), "CreateDirStat");
+  const string dir = TmpDir("CreateDirStat");
   TF_EXPECT_OK(hdfs.CreateDir(dir));
   FileStatistics stat;
   TF_EXPECT_OK(hdfs.Stat(dir, &stat));
@@ -148,7 +150,7 @@ TEST_F(HadoopFileSystemTest, CreateDirStat) {
 }
 
 TEST_F(HadoopFileSystemTest, DeleteDir) {
-  const string dir = "file://" + io::JoinPath(testing::TmpDir(), "DeleteDir");
+  const string dir = TmpDir("DeleteDir");
   EXPECT_FALSE(hdfs.DeleteDir(dir).ok());
   TF_EXPECT_OK(hdfs.CreateDir(dir));
   TF_EXPECT_OK(hdfs.DeleteDir(dir));
@@ -157,10 +159,8 @@ TEST_F(HadoopFileSystemTest, DeleteDir) {
 }
 
 TEST_F(HadoopFileSystemTest, RenameFile) {
-  const string fname1 =
-      "file://" + io::JoinPath(testing::TmpDir(), "RenameFile1");
-  const string fname2 =
-      "file://" + io::JoinPath(testing::TmpDir(), "RenameFile2");
+  const string fname1 = TmpDir("RenameFile1");
+  const string fname2 = TmpDir("RenameFile2");
   TF_ASSERT_OK(WriteString(fname1, "test"));
   TF_EXPECT_OK(hdfs.RenameFile(fname1, fname2));
   string content;
@@ -169,10 +169,8 @@ TEST_F(HadoopFileSystemTest, RenameFile) {
 }
 
 TEST_F(HadoopFileSystemTest, RenameFile_Overwrite) {
-  const string fname1 =
-      "file://" + io::JoinPath(testing::TmpDir(), "RenameFile1");
-  const string fname2 =
-      "file://" + io::JoinPath(testing::TmpDir(), "RenameFile2");
+  const string fname1 = TmpDir("RenameFile1");
+  const string fname2 = TmpDir("RenameFile2");
 
   TF_ASSERT_OK(WriteString(fname2, "test"));
   TF_EXPECT_OK(hdfs.FileExists(fname2));
@@ -185,7 +183,7 @@ TEST_F(HadoopFileSystemTest, RenameFile_Overwrite) {
 }
 
 TEST_F(HadoopFileSystemTest, StatFile) {
-  const string fname = "file://" + io::JoinPath(testing::TmpDir(), "StatFile");
+  const string fname = TmpDir("StatFile");
   TF_ASSERT_OK(WriteString(fname, "test"));
   FileStatistics stat;
   TF_EXPECT_OK(hdfs.Stat(fname, &stat));