trial_dir.py 2.68 KB
Newer Older
1
# Copyright (c) 2011 The Chromium Authors. All rights reserved.
2 3 4 5 6 7 8 9 10
# 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 sys
import tempfile
11 12

from testing_support import auto_stub
13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45

import gclient_utils


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)
46
    gclient_utils.rmtree(self.root_dir)
47 48 49 50 51 52
    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)
53
      gclient_utils.rmtree(self.root_dir)
54 55 56 57 58 59 60 61 62
    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)
63
      gclient_utils.rmtree(TrialDir.TRIAL_ROOT)
64 65 66 67
    else:
      logging.error('Leaking %s' % TrialDir.TRIAL_ROOT)


68
class TrialDirMixIn(object):
69 70 71 72 73 74 75 76 77 78 79 80 81
  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


82
class TestCase(auto_stub.TestCase, TrialDirMixIn):
83 84
  """Base unittest class that cleans off a trial directory in tearDown()."""
  def setUp(self):
85
    auto_stub.TestCase.setUp(self)
86 87 88 89
    TrialDirMixIn.setUp(self)

  def tearDown(self):
    TrialDirMixIn.tearDown(self)
90
    auto_stub.TestCase.tearDown(self)
91 92


93 94 95 96 97
if '-l' in sys.argv:
  # See SHOULD_LEAK definition in TrialDir for its purpose.
  TrialDir.SHOULD_LEAK = True
  print 'Leaking!'
  sys.argv.remove('-l')