Added TempDir

expect_tests/pipeline.py

 # Copyright 2014 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 import Queue
 import glob
 import logging
 import multiprocessing
 import os
 import re
 import shutil
 import signal
 import sys
 import tempfile
 import traceback
 
 from cStringIO import StringIO
 
+from expect_tests.tempdir import TempDir
 from expect_tests.type_definitions import (
     Test, UnknownError, TestError, NoMatchingTestsError, MultiTest,
     Result, ResultStageAbort)
 
 from expect_tests import util
 
 
-OLD_TEMP_DIR = None
-TEMP_PID_FILE = '.expect_tests_pidfile'
-
-
 class ResetableStringIO(object):
   def __init__(self):
     self._stream = StringIO()
 
   def reset(self):
     self._stream = StringIO()
 
   def __getattr__(self, key):
     return getattr(self._stream, key)
 
 
-def set_temp_dir():
-  """Provides an automatically-cleaned temporary directory base for tests."""
-  global OLD_TEMP_DIR
-
-  if OLD_TEMP_DIR:
-    return
-
-  suffix = '.expect_tests_temp'
-  tempdir = tempfile.mkdtemp(suffix)
-  for p in glob.glob(os.path.join(os.path.dirname(tempdir), '*'+suffix)):
-    if p == tempdir:
-      continue
-
-    pfile = os.path.join(p, TEMP_PID_FILE)
-    if os.path.exists(pfile):
-      with open(pfile, 'rb') as f:
-        try:
-          os.kill(int(f.read()), 0)
-          continue
-        except (OSError, TypeError):
-          pass
-
-    shutil.rmtree(p, ignore_errors=True)
-
-  with open(os.path.join(tempdir, TEMP_PID_FILE), 'wb') as f:
-    f.write(str(os.getpid()))
-
-  OLD_TEMP_DIR = tempfile.tempdir
-  tempfile.tempdir = tempdir
-
-
-def clear_temp_dir():
-  global OLD_TEMP_DIR
-  if OLD_TEMP_DIR:
-    # Try to nuke the pidfile first
-    pfile = os.path.join(tempfile.tempdir, TEMP_PID_FILE)
-    try:
-      os.unlink(pfile)
-    except OSError as e:
-      print >> sys.stderr, "Error removing %r: %s" % (pfile, e)
-    shutil.rmtree(tempfile.tempdir, ignore_errors=True)
-    tempfile.temdir = OLD_TEMP_DIR
-    OLD_TEMP_DIR = None
-
-
 def gen_loop_process(gens, test_queue, result_queue, opts, kill_switch,
                      cover_ctx, temp_dir):
   """Generate `Test`s from |gens|, and feed them into |test_queue|.
 
   Non-Test instances will be translated into `UnknownError` objects.
 
   On completion, feed |opts.jobs| None objects into |test_queue|.
 
   @param gens: list of generators yielding Test() instances.
   @type test_queue: multiprocessing.Queue()
@@ skipping 159 matching lines @@
     kill_switch.set()
     # Reset the signal to DFL so that double ctrl-C kills us for sure.
     signal.signal(signal.SIGINT, signal.SIG_DFL)
     signal.signal(signal.SIGTERM, signal.SIG_DFL)
   signal.signal(signal.SIGINT, handle_killswitch)
   signal.signal(signal.SIGTERM, handle_killswitch)
 
   test_queue = multiprocessing.Queue()
   result_queue = multiprocessing.Queue()
 
-  set_temp_dir()
+  with TempDir() as temp_dir:
+    test_gen_args = (
+        test_gens, test_queue, result_queue, opts, kill_switch, cover_ctx,
+        temp_dir
+    )
 
-  test_gen_args = (
-      test_gens, test_queue, result_queue, opts, kill_switch, cover_ctx,
-      tempfile.tempdir
-  )
+    procs = []
+    if opts.handler.SKIP_RUNLOOP:
+      gen_loop_process(*test_gen_args)
+    else:
+      procs = [multiprocessing.Process(
+          target=gen_loop_process, args=test_gen_args)]
 
-  procs = []
-  if opts.handler.SKIP_RUNLOOP:
-    gen_loop_process(*test_gen_args)
-  else:
-    procs = [multiprocessing.Process(
-        target=gen_loop_process, args=test_gen_args)]
+      procs += [
+          multiprocessing.Process(
+              target=run_loop_process, args=(
+                  test_queue, result_queue, opts, kill_switch, cover_ctx,
+                  temp_dir))
+          for _ in xrange(opts.jobs)
+      ]
 
-    procs += [
-        multiprocessing.Process(
-            target=run_loop_process, args=(
-                test_queue, result_queue, opts, kill_switch, cover_ctx,
-                tempfile.tempdir))
-        for _ in xrange(opts.jobs)
-    ]
+      for p in procs:
+        p.daemon = True
+        p.start()
 
-    for p in procs:
-      p.daemon = True
-      p.start()
+    error = False
 
-  error = False
+    try:
+      def generate_objects():
+        while not kill_switch.is_set():
+          while not kill_switch.is_set():
+            try:
+              yield result_queue.get(timeout=0.1)
+            except Queue.Empty:
+              break
 
-  try:
-    def generate_objects():
-      while not kill_switch.is_set():
+          if not any(p.is_alive() for p in procs):
+            break
+
+        # Get everything still in the queue. Still need timeout, but since
+        # nothing is going to be adding stuff to the queue, use a very short
+        # timeout.
         while not kill_switch.is_set():
           try:
-            yield result_queue.get(timeout=0.1)
+            yield result_queue.get(timeout=0.00001)
           except Queue.Empty:
             break
 
-        if not any(p.is_alive() for p in procs):
-          break
-
-      # Get everything still in the queue. Still need timeout, but since nothing
-      # is going to be adding stuff to the queue, use a very short timeout.
-      while not kill_switch.is_set():
-        try:
-          yield result_queue.get(timeout=0.00001)
-        except Queue.Empty:
-          break
-
-      if kill_switch.is_set():
-        raise ResultStageAbort()
-    error = opts.handler.result_stage_loop(opts, generate_objects())
-  except ResultStageAbort:
-    pass
-  finally:
-    clear_temp_dir()
+        if kill_switch.is_set():
+          raise ResultStageAbort()
+      error = opts.handler.result_stage_loop(opts, generate_objects())
+    except ResultStageAbort:
+      pass
 
   for p in procs:
     p.join()
 
   if not kill_switch.is_set() and not result_queue.empty():
     error = True
 
   return error, kill_switch.is_set()