testing_support/trial_dir.py

# Copyright (c) 2011 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 atexit
import logging
import os
import stat
import subprocess
import sys
import tempfile
import time

from testing_support import auto_stub


def rmtree(path):
  """shutil.rmtree() on steroids.

  Recursively removes a directory, even if it's marked read-only.

  shutil.rmtree() doesn't work on Windows if any of the files or directories
  are read-only, which svn repositories and some .svn files are.  We need to
  be able to force the files to be writable (i.e., deletable) as we traverse
  the tree.

  Even with all this, Windows still sometimes fails to delete a file, citing
  a permission error (maybe something to do with antivirus scans or disk
  indexing).  The best suggestion any of the user forums had was to wait a
  bit and try again, so we do that too.  It's hand-waving, but sometimes it
  works. :/

  On POSIX systems, things are a little bit simpler.  The modes of the files
  to be deleted doesn't matter, only the modes of the directories containing
  them are significant.  As the directory tree is traversed, each directory
  has its mode set appropriately before descending into it.  This should
  result in the entire tree being removed, with the possible exception of
  *path itself, because nothing attempts to change the mode of its parent.
  Doing so would be hazardous, as it's not a directory slated for removal.
  In the ordinary case, this is not a problem: for our purposes, the user
  will never lack write permission on *path's parent.
  """
  if not os.path.exists(path):
    return

  if os.path.islink(path) or not os.path.isdir(path):
    raise ValueError('Called rmtree(%s) in non-directory' % path)

  if sys.platform == 'win32':
    # Give up and use cmd.exe's rd command.
    path = os.path.normcase(path)
    for _ in xrange(3):
      exitcode = subprocess.call(['cmd.exe', '/c', 'rd', '/q', '/s', path])
      if exitcode == 0:
        return
      else:
        print >> sys.stderr, 'rd exited with code %d' % exitcode
      time.sleep(3)
    raise Exception('Failed to remove path %s' % path)

  # On POSIX systems, we need the x-bit set on the directory to access it,
  # the r-bit to see its contents, and the w-bit to remove files from it.
  # The actual modes of the files within the directory is irrelevant.
  os.chmod(path, stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR)

  def remove(func, subpath):
    func(subpath)

  for fn in os.listdir(path):
    # If fullpath is a symbolic link that points to a directory, isdir will
    # be True, but we don't want to descend into that as a directory, we just
    # want to remove the link.  Check islink and treat links as ordinary files
    # would be treated regardless of what they reference.
    fullpath = os.path.join(path, fn)
    if os.path.islink(fullpath) or not os.path.isdir(fullpath):
      remove(os.remove, fullpath)
    else:
      # Recurse.
      rmtree(fullpath)

  remove(os.rmdir, path)


class TrialDir(object):
  """Manages a temporary directory.

  On first object creation, TrialDir.TRIAL_ROOT will be set to a new temporary
  directory created in /tmp or the equivalent. It will be deleted on process
  exit unless TrialDir.SHOULD_LEAK is set to True.
  """
  # When SHOULD_LEAK is set to True, temporary directories created while the
  # tests are running aren't deleted at the end of the tests. Expect failures
  # when running more than one test due to inter-test side-effects. Helps with
  # debugging.
  SHOULD_LEAK = False

  # Main root directory.
  TRIAL_ROOT = None

  def __init__(self, subdir, leak=False):
    self.leak = self.SHOULD_LEAK or leak
    self.subdir = subdir
    self.root_dir = None

  def set_up(self):
    """All late initialization comes here."""
    # You can override self.TRIAL_ROOT.
    if not self.TRIAL_ROOT:
      # Was not yet initialized.
      TrialDir.TRIAL_ROOT = os.path.realpath(tempfile.mkdtemp(prefix='trial'))
      atexit.register(self._clean)
    self.root_dir = os.path.join(TrialDir.TRIAL_ROOT, self.subdir)
    rmtree(self.root_dir)
    os.makedirs(self.root_dir)

  def tear_down(self):
    """Cleans the trial subdirectory for this instance."""
    if not self.leak:
      logging.debug('Removing %s' % self.root_dir)
      rmtree(self.root_dir)
    else:
      logging.error('Leaking %s' % self.root_dir)
    self.root_dir = None

  @staticmethod
  def _clean():
    """Cleans the root trial directory."""
    if not TrialDir.SHOULD_LEAK:
      logging.debug('Removing %s' % TrialDir.TRIAL_ROOT)
      rmtree(TrialDir.TRIAL_ROOT)
    else:
      logging.error('Leaking %s' % TrialDir.TRIAL_ROOT)


class TrialDirMixIn(object):
  def setUp(self):
    # Create a specific directory just for the test.
    self.trial = TrialDir(self.id())
    self.trial.set_up()

  def tearDown(self):
    self.trial.tear_down()

  @property
  def root_dir(self):
    return self.trial.root_dir


class TestCase(auto_stub.TestCase, TrialDirMixIn):
  """Base unittest class that cleans off a trial directory in tearDown()."""
  def setUp(self):
    TrialDirMixIn.setUp(self)
    auto_stub.TestCase.setUp(self)

  def tearDown(self):
    auto_stub.TestCase.tearDown(self)
    TrialDirMixIn.tearDown(self)