chrome/test/pyautolib/pyauto_utils.py

   1 # Copyright (c) 2012 The Chromium Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4
   5 """Utilities for PyAuto."""
   6
   7 import httplib
   8 import logging
   9 import os
  10 import shutil
  11 import socket
  12 import sys
  13 import tempfile
  14 import unittest
  15 import urlparse
  16 import zipfile
  17
  18
  19 class ExistingPathReplacer(object):
  20   """Facilitates backing up a given path (file or dir)..
  21
  22   Often you want to manipulate a directory or file for testing but don't want to
  23   meddle with the existing contents.  This class lets you make a backup, and
  24   reinstate the backup when done.  A backup is made in an adjacent directory,
  25   so you need to make sure you have write permissions to the parent directory.
  26
  27   Works seemlessly in cases where the requested path already exists, or not.
  28
  29   Automatically reinstates the backed up path (if any) when object is deleted.
  30   """
  31   _path = ''
  32   _backup_dir = None  # dir to which existing content is backed up
  33   _backup_basename = ''
  34
  35   def __init__(self, path, path_type='dir'):
  36     """Initialize the object, making backups if necessary.
  37
  38     Args:
  39       path: the requested path to file or directory
  40       path_type: path type. Options: 'file', 'dir'. Default: 'dir'
  41     """
  42     assert path_type in ('file', 'dir'), 'Invalid path_type: %s' % path_type
  43     self._path_type = path_type
  44     self._path = path
  45     if os.path.exists(self._path):
  46       if 'dir' == self._path_type:
  47         assert os.path.isdir(self._path), '%s is not a directory' % self._path
  48       else:
  49         assert os.path.isfile(self._path), '%s is not a file' % self._path
  50       # take a backup
  51       self._backup_basename = os.path.basename(self._path)
  52       self._backup_dir = tempfile.mkdtemp(dir=os.path.dirname(self._path),
  53                                           prefix='bkp-' + self._backup_basename)
  54       logging.info('Backing up %s in %s' % (self._path, self._backup_dir))
  55       shutil.move(self._path,
  56                   os.path.join(self._backup_dir, self._backup_basename))
  57     self._CreateRequestedPath()
  58
  59   def __del__(self):
  60     """Cleanup. Reinstate backup."""
  61     self._CleanupRequestedPath()
  62     if self._backup_dir:  # Reinstate, if backed up.
  63       from_path = os.path.join(self._backup_dir, self._backup_basename)
  64       logging.info('Reinstating backup from %s to %s' % (from_path, self._path))
  65       shutil.move(from_path, self._path)
  66     self._RemoveBackupDir()
  67
  68   def _CreateRequestedPath(self):
  69     # Create intermediate dirs if needed.
  70     if not os.path.exists(os.path.dirname(self._path)):
  71       os.makedirs(os.path.dirname(self._path))
  72     if 'dir' == self._path_type:
  73       os.mkdir(self._path)
  74     else:
  75       open(self._path, 'w').close()
  76
  77   def _CleanupRequestedPath(self):
  78     if os.path.exists(self._path):
  79       if os.path.isdir(self._path):
  80         shutil.rmtree(self._path, ignore_errors=True)
  81       else:
  82         os.remove(self._path)
  83
  84   def _RemoveBackupDir(self):
  85     if self._backup_dir and os.path.isdir(self._backup_dir):
  86       shutil.rmtree(self._backup_dir, ignore_errors=True)
  87
  88
  89 def RemovePath(path):
  90   """Remove the given path (file or dir)."""
  91   if os.path.isdir(path):
  92     shutil.rmtree(path, ignore_errors=True)
  93     return
  94   try:
  95     os.remove(path)
  96   except OSError:
  97     pass
  98
  99
 100 def UnzipFilenameToDir(filename, dir):
 101   """Unzip |filename| to directory |dir|.
 102
 103   This works with as low as python2.4 (used on win).
 104   """
 105   zf = zipfile.ZipFile(filename)
 106   pushd = os.getcwd()
 107   if not os.path.isdir(dir):
 108     os.mkdir(dir)
 109   os.chdir(dir)
 110   # Extract files.
 111   for info in zf.infolist():
 112     name = info.filename
 113     if name.endswith('/'):  # dir
 114       if not os.path.isdir(name):
 115         os.makedirs(name)
 116     else:  # file
 117       dir = os.path.dirname(name)
 118       if not os.path.isdir(dir):
 119         os.makedirs(dir)
 120       out = open(name, 'wb')
 121       out.write(zf.read(name))
 122       out.close()
 123     # Set permissions. Permission info in external_attr is shifted 16 bits.
 124     os.chmod(name, info.external_attr >> 16L)
 125   os.chdir(pushd)
 126
 127
 128 def GetCurrentPlatform():
 129   """Get a string representation for the current platform.
 130
 131   Returns:
 132     'mac', 'win' or 'linux'
 133   """
 134   if sys.platform == 'darwin':
 135     return 'mac'
 136   if sys.platform == 'win32':
 137     return 'win'
 138   if sys.platform.startswith('linux'):
 139     return 'linux'
 140   raise RuntimeError('Unknown platform')
 141
 142
 143 def PrintPerfResult(graph_name, series_name, data_point, units,
 144                     show_on_waterfall=False):
 145   """Prints a line to stdout that is specially formatted for the perf bots.
 146
 147   Args:
 148     graph_name: String name for the graph on which to plot the data.
 149     series_name: String name for the series (line on the graph) associated with
 150                  the data.  This is also the string displayed on the waterfall
 151                  if |show_on_waterfall| is True.
 152     data_point: Numeric data value to plot on the graph for the current build.
 153                 This can be a single value or an array of values.  If an array,
 154                 the graph will plot the average of the values, along with error
 155                 bars.
 156     units: The string unit of measurement for the given |data_point|.
 157     show_on_waterfall: Whether or not to display this result directly on the
 158                        buildbot waterfall itself (in the buildbot step running
 159                        this test on the waterfall page, not the stdio page).
 160   """
 161   waterfall_indicator = ['', '*'][show_on_waterfall]
 162   print '%sRESULT %s: %s= %s %s' % (
 163       waterfall_indicator, graph_name, series_name,
 164       str(data_point).replace(' ', ''), units)
 165   sys.stdout.flush()
 166
 167
 168 def Shard(ilist, shard_index, num_shards):
 169   """Shard a given list and return the group at index |shard_index|.
 170
 171   Args:
 172     ilist: input list
 173     shard_index: 0-based sharding index
 174     num_shards: shard count
 175   """
 176   chunk_size = len(ilist) / num_shards
 177   chunk_start = shard_index * chunk_size
 178   if shard_index == num_shards - 1:  # Exhaust the remainder in the last shard.
 179     chunk_end = len(ilist)
 180   else:
 181     chunk_end = chunk_start + chunk_size
 182   return ilist[chunk_start:chunk_end]
 183
 184
 185 def WaitForDomElement(pyauto, driver, xpath):
 186   """Wait for the UI element to appear.
 187
 188   Args:
 189     pyauto: an instance of pyauto.PyUITest.
 190     driver: an instance of chrome driver or a web element.
 191     xpath: the xpath of the element to wait for.
 192
 193   Returns:
 194     The element if it is found.
 195     NoSuchElementException if it is not found.
 196   """
 197   pyauto.WaitUntil(lambda: len(driver.find_elements_by_xpath(xpath)) > 0)
 198   return driver.find_element_by_xpath(xpath)
 199
 200
 201 def DoesUrlExist(url):
 202   """Determines whether a resource exists at the given URL.
 203
 204   Args:
 205     url: URL to be verified.
 206
 207   Returns:
 208     True if url exists, otherwise False.
 209   """
 210   parsed = urlparse.urlparse(url)
 211   try:
 212     conn = httplib.HTTPConnection(parsed.netloc)
 213     conn.request('HEAD', parsed.path)
 214     response = conn.getresponse()
 215   except (socket.gaierror, socket.error):
 216     return False
 217   finally:
 218     conn.close()
 219   # Follow both permanent (301) and temporary (302) redirects.
 220   if response.status == 302 or response.status == 301:
 221     return DoesUrlExist(response.getheader('location'))
 222   return response.status == 200
 223
 224
 225 class _GTestTextTestResult(unittest._TextTestResult):
 226   """A test result class that can print formatted text results to a stream.
 227
 228   Results printed in conformance with gtest output format, like:
 229   [ RUN        ] autofill.AutofillTest.testAutofillInvalid: "test desc."
 230   [         OK ] autofill.AutofillTest.testAutofillInvalid
 231   [ RUN        ] autofill.AutofillTest.testFillProfile: "test desc."
 232   [         OK ] autofill.AutofillTest.testFillProfile
 233   [ RUN        ] autofill.AutofillTest.testFillProfileCrazyCharacters: "Test."
 234   [         OK ] autofill.AutofillTest.testFillProfileCrazyCharacters
 235   """
 236
 237   def __init__(self, stream, descriptions, verbosity):
 238     unittest._TextTestResult.__init__(self, stream, descriptions, verbosity)
 239
 240   def _GetTestURI(self, test):
 241     if sys.version_info[:2] <= (2, 4):
 242       return '%s.%s' % (unittest._strclass(test.__class__),
 243                         test._TestCase__testMethodName)
 244     return '%s.%s' % (unittest._strclass(test.__class__), test._testMethodName)
 245
 246   def getDescription(self, test):
 247     return '%s: "%s"' % (self._GetTestURI(test), test.shortDescription())
 248
 249   def startTest(self, test):
 250     unittest.TestResult.startTest(self, test)
 251     self.stream.writeln('[ RUN        ] %s' % self.getDescription(test))
 252
 253   def addSuccess(self, test):
 254     unittest.TestResult.addSuccess(self, test)
 255     self.stream.writeln('[         OK ] %s' % self._GetTestURI(test))
 256
 257   def addError(self, test, err):
 258     unittest.TestResult.addError(self, test, err)
 259     self.stream.writeln('[      ERROR ] %s' % self._GetTestURI(test))
 260
 261   def addFailure(self, test, err):
 262     unittest.TestResult.addFailure(self, test, err)
 263     self.stream.writeln('[     FAILED ] %s' % self._GetTestURI(test))
 264
 265
 266 class GTestTextTestRunner(unittest.TextTestRunner):
 267   """Test Runner for displaying test results in textual format.
 268
 269   Results are displayed in conformance with gtest output.
 270   """
 271
 272   def __init__(self, verbosity=1):
 273     unittest.TextTestRunner.__init__(self, stream=sys.stderr,
 274                                      verbosity=verbosity)
 275
 276   def _makeResult(self):
 277     return _GTestTextTestResult(self.stream, self.descriptions, self.verbosity)