Bazel, testing: improve test_base.py

Changes to test_base.py:
- use tempfile.TemporaryFile to create Bazel's
  stdout/stderr files
- added a new RunProgram method
- setUp now creates a unique temp directory for
  each test case

Change-Id: I5b851fee826e73044fa3992a5981855982527ae3
PiperOrigin-RevId: 160505154

1 files changed, 47 insertions, 24 deletions

diff --git a/src/test/py/bazel/test_base.py b/src/test/py/bazel/test_base.py
index e782f9d3b7..d0a9ba7303 100644
--- a/src/test/py/bazel/test_base.py
+++ b/src/test/py/bazel/test_base.py
@@ -15,6 +15,7 @@
 import os
 import subprocess
 import sys
+import tempfile
 import unittest
 
 
@@ -37,23 +38,21 @@ class EnvVarUndefinedError(Error):
 
 class TestBase(unittest.TestCase):
 
-  _stderr = None
-  _stdout = None
   _runfiles = None
   _temp = None
   _tests_root = None
+  _test_cwd = None
 
   def setUp(self):
     unittest.TestCase.setUp(self)
     if self._runfiles is None:
       self._runfiles = TestBase._LoadRunfiles()
-    test_tmpdir = TestBase.GetEnv('TEST_TMPDIR')
-    self._stdout = os.path.join(test_tmpdir, 'bazel.stdout')
-    self._stderr = os.path.join(test_tmpdir, 'bazel.stderr')
-    self._temp = os.path.join(test_tmpdir, 'tmp')
-    self._tests_root = os.path.join(test_tmpdir, 'tests_root')
-    os.mkdir(self._tests_root)
-    os.chdir(self._tests_root)
+    test_tmpdir = TestBase._CreateDirs(TestBase.GetEnv('TEST_TMPDIR'))
+    self._tests_root = TestBase._CreateDirs(
+        os.path.join(test_tmpdir, 'tests_root'))
+    self._temp = TestBase._CreateDirs(os.path.join(test_tmpdir, 'tmp'))
+    self._test_cwd = tempfile.mkdtemp(dir=self._tests_root)
+    os.chdir(self._test_cwd)
 
   def AssertExitCode(self, actual_exit_code, expected_exit_code, stderr_lines):
     """Assert that `actual_exit_code` == `expected_exit_code`."""
@@ -168,27 +167,42 @@ class TestBase(unittest.TestCase):
     Returns:
       (int, [string], [string]) tuple: exit code, stdout lines, stderr lines
     """
-    with open(self._stdout, 'w') as stdout:
-      with open(self._stderr, 'w') as stderr:
+    return self.RunProgram([
+        self.Rlocation('io_bazel/src/bazel'),
+        '--bazelrc=/dev/null',
+        '--nomaster_bazelrc',
+    ] + args, env_remove)
+
+  def RunProgram(self, args, env_remove=None):
+    """Runs a program (args[0]), waits for it to exit.
+
+    Args:
+      args: [string]; the args to run; args[0] should be the program itself
+      env_remove: set(string); optional; environment variables to NOT pass to
+        the program
+    Returns:
+      (int, [string], [string]) tuple: exit code, stdout lines, stderr lines
+    """
+    with tempfile.TemporaryFile(dir=self._test_cwd) as stdout:
+      with tempfile.TemporaryFile(dir=self._test_cwd) as stderr:
         proc = subprocess.Popen(
-            [
-                self.Rlocation('io_bazel/src/bazel'), '--bazelrc=/dev/null',
-                '--nomaster_bazelrc'
-            ] + args,
+            args,
             stdout=stdout,
             stderr=stderr,
-            cwd=self._tests_root,
-            env=self._BazelEnvMap(env_remove))
+            cwd=self._test_cwd,
+            env=self._EnvMap(env_remove))
         exit_code = proc.wait()
 
-    with open(self._stdout, 'r') as f:
-      stdout = [l.strip() for l in f.readlines()]
-    with open(self._stderr, 'r') as f:
-      stderr = [l.strip() for l in f.readlines()]
-    return exit_code, stdout, stderr
+        stdout.seek(0)
+        stdout_lines = [l.strip() for l in stdout.readlines()]
+
+        stderr.seek(0)
+        stderr_lines = [l.strip() for l in stderr.readlines()]
+
+        return exit_code, stdout_lines, stderr_lines
 
-  def _BazelEnvMap(self, env_remove=None):
-    """Returns the environment variable map to run Bazel."""
+  def _EnvMap(self, env_remove=None):
+    """Returns the environment variable map to run Bazel or other programs."""
     if TestBase.IsWindows():
       result = []
       if sys.version_info.major == 3:
@@ -253,3 +267,12 @@ class TestBase(unittest.TestCase):
         if len(tokens) == 2:
           result[tokens[0]] = tokens[1]
     return result
+
+  @staticmethod
+  def _CreateDirs(path):
+    if not os.path.exists(path):
+      os.makedirs(path)
+    elif not os.path.isdir(path):
+      os.remove(path)
+      os.makedirs(path)
+    return path