Add trial_dir to manage automatic trial directory creation and deletion.

tests/trial_dir.py

Index: tests/trial_dir.py
diff --git a/tests/trial_dir.py b/tests/trial_dir.py
new file mode 100755
index 0000000000000000000000000000000000000000..d75c5d71b987683ceeb63514d0f490d5adf979f8
--- /dev/null
+++ b/tests/trial_dir.py
@@ -0,0 +1,87 @@
+#!/usr/bin/python
+# Copyright (c) 2010 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 sys
+import tempfile
+import unittest
+
+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)
+    gclient_utils.RemoveDirectory(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)
+      gclient_utils.RemoveDirectory(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)
+      gclient_utils.RemoveDirectory(TrialDir.TRIAL_ROOT)
+    else:
+      logging.error('Leaking %s' % TrialDir.TRIAL_ROOT)
+
+
+class TestCase(unittest.TestCase):
+  """Base unittest class that cleans off a trial directory in tearDown()."""
+  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
+
+
+if '-l' in sys.argv:
+  # See SHOULD_LEAK definition in TrialDir for its purpose.
+  TrialDir.SHOULD_LEAK = True
+  print 'Leaking!'
+  sys.argv.remove('-l')