Add some desired options to the au test harness.

bin/cros_au_test_harness.py

 #!/usr/bin/python
 
 # Copyright (c) 2010 The Chromium OS 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 optparse
 import os
 import re
 import sys
@@ skipping 24 matching lines @@
 class UpdateException(Exception):
   """Exception thrown when UpdateImage or UpdateUsingPayload fail"""
   def __init__(self, code, stdout):
     self.code = code
     self.stdout = stdout
 
 class AUTest(object):
   """Abstract interface that defines an Auto Update test."""
   source_image = ''
   use_delta_updates = False
+  verbose = False
 
   def setUp(self):
     unittest.TestCase.setUp(self)
     # Set these up as they are used often.
     self.crosutils = os.path.join(os.path.dirname(__file__), '..')
     self.crosutilsbin = os.path.join(os.path.dirname(__file__))
     self.download_folder = os.path.join(self.crosutils, 'latest_download')
     if not os.path.exists(self.download_folder):
       os.makedirs(self.download_folder)
 
@@ skipping 18 matching lines @@
         break
 
     return int(percent_passed)
 
   # TODO(sosa) - Remove try and convert function to DeltaUpdateImage().
   def TryDeltaAndFallbackToFull(self, src_image, image, stateful_change='old'):
     """Tries the delta update first if set and falls back to full update."""
     if self.use_delta_updates:
       try:
         self.source_image = src_image
-        self.UpdateImage(image)
+        self._UpdateImageReportError(image)
       except:
         Warning('Delta update failed, disabling delta updates and retrying.')
         self.use_delta_updates = False
         self.source_image = ''
         self._UpdateImageReportError(image)
     else:
       self._UpdateImageReportError(image)
 
   def _UpdateImageReportError(self, image_path, stateful_change='old'):
     """Calls UpdateImage and reports any error to the console.
 
        Still throws the exception.
     """
     try:
       self.UpdateImage(image_path, stateful_change)
     except UpdateException as err:
       # If the update fails, print it out
       Warning(err.stdout)
       raise
 
   def _AttemptUpdateWithPayloadExpectedFailure(self, payload, expected_msg):
     # This update is expected to fail...
     try:
       self.UpdateUsingPayload(payload)
     except UpdateException as err:
       # Will raise ValueError if expected is not found.
       if re.search(re.escape(expected_msg), err.stdout, re.MULTILINE):
         return
 
-      Warning("Didn't find '%s' in:" % expected_msg)
-      Warning(err.stdout)
-      self.fail('We managed to update when failure was expected')
+    Warning("Didn't find '%s' in:" % expected_msg)
+    Warning(err.stdout)
+    self.fail('We managed to update when failure was expected')
 
   def PrepareBase(self, image_path):
     """Prepares target with base_image_path."""
     pass
 
   def UpdateImage(self, image_path, stateful_change='old'):
     """Updates target with the image given by the image_path.
 
     Args:
       image_path:  Path to the image to update with.  This image must be a test
@@ skipping 36 matching lines @@
 
     Args:
       unittest: Handle to the unittest.
       output: stdout from a test run.
       percent_required_to_pass: percentage required to pass.  This should be
         fall between 0-100.
     Returns:
       percent that passed.
     """
     Info('Output from VerifyImage():')
-    print output
+    print >> sys.stderr, output
+    sys.stderr.flush()
     percent_passed = self.ParseGenerateTestReportOutput(output)
     Info('Percent passed: %d vs. Percent required: %d' % (
         percent_passed, percent_required_to_pass))
     unittest.assertTrue(percent_passed >=
                         percent_required_to_pass)
     return percent_passed
 
   def testFullUpdateKeepStateful(self):
     """Tests if we can update normally.
 
@@ skipping 47 matching lines @@
 
     # Image can be updated at:
     # ~chrome-eng/chromeos/localmirror/autest-images
     url = 'http://gsdview.appspot.com/chromeos-localmirror/' \
           'autest-images/truncated_image.gz'
     payload = os.path.join(self.download_folder, 'truncated_image.gz')
 
     # Read from the URL and write to the local file
     urllib.urlretrieve(url, payload)
 
-    expected_msg='download_hash_data == update_check_response_hash failed'
+    expected_msg = 'download_hash_data == update_check_response_hash failed'
     self._AttemptUpdateWithPayloadExpectedFailure(payload, expected_msg)
 
   def testCorruptedUpdate(self):
     """Tests what happens if we attempt to update with a corrupted payload."""
     # Preload with the version we are trying to test.
     self.PrepareBase(image_path=target_image_path)
 
     # Image can be updated at:
     # ~chrome-eng/chromeos/localmirror/autest-images
     url = 'http://gsdview.appspot.com/chromeos-localmirror/' \
           'autest-images/corrupted_image.gz'
     payload = os.path.join(self.download_folder, 'corrupted.gz')
 
     # Read from the URL and write to the local file
     urllib.urlretrieve(url, payload)
 
     # This update is expected to fail...
-    expected_msg='zlib inflate() error:-3'
+    expected_msg = 'zlib inflate() error:-3'
     self._AttemptUpdateWithPayloadExpectedFailure(payload, expected_msg)
 
 class RealAUTest(unittest.TestCase, AUTest):
   """Test harness for updating real images."""
 
   def setUp(self):
     AUTest.setUp(self)
 
   def PrepareBase(self, image_path):
     """Auto-update to base image to prepare for test."""
     self._UpdateImageReportError(image_path)
 
   def UpdateImage(self, image_path, stateful_change='old'):
     """Updates a remote image using image_to_live.sh."""
     stateful_change_flag = self.GetStatefulChangeFlag(stateful_change)
+    cmd = ['%s/image_to_live.sh' % self.crosutils,
+           '--image=%s' % image_path,
+           '--remote=%s' % remote,
+           stateful_change_flag,
+           '--verify',
+           '--src_image=%s' % self.source_image
+          ]
 
-    (code, stdout, stderr) = RunCommandCaptureOutput([
-        '%s/image_to_live.sh' % self.crosutils,
-        '--image=%s' % image_path,
-        '--remote=%s' % remote,
-        stateful_change_flag,
-        '--verify',
-        '--src_image=%s' % self.source_image
-        ])
-
-    if code != 0:
-      raise UpdateException(code, stdout)
+    if self.verbose:
+      try:
+        RunCommand(cmd)
+      except Exception, e:
+        raise UpdateException(1, e.message)
+    else:
+      (code, stdout, stderr) = RunCommandCaptureOutput(cmd)
+      if code != 0:
+        raise UpdateException(code, stdout)
 
   def UpdateUsingPayload(self, update_path, stateful_change='old'):
     """Updates a remote image using image_to_live.sh."""
     stateful_change_flag = self.GetStatefulChangeFlag(stateful_change)
+    cmd = ['%s/image_to_live.sh' % self.crosutils,
+           '--payload=%s' % update_path,
+           '--remote=%s' % remote,
+           stateful_change_flag,
+           '--verify',
+          ]
 
-    (code, stdout, stderr) = RunCommandCaptureOutput([
-        '%s/image_to_live.sh' % self.crosutils,
-        '--payload=%s' % update_path,
-        '--remote=%s' % remote,
-        stateful_change_flag,
-        '--verify',
-        ])
-
-    if code != 0:
-      raise UpdateException(code, stdout)
+    if self.verbose:
+      try:
+        RunCommand(cmd)
+      except Exception, e:
+        raise UpdateException(1, e.message)
+    else:
+      (code, stdout, stderr) = RunCommandCaptureOutput(cmd)
+      if code != 0:
+        raise UpdateException(code, stdout)
 
   def VerifyImage(self, percent_required_to_pass):
     """Verifies an image using run_remote_tests.sh with verification suite."""
     output = RunCommand([
         '%s/run_remote_tests.sh' % self.crosutils,
         '--remote=%s' % remote,
         _VERIFY_SUITE,
        ], error_ok=True, enter_chroot=False, redirect_stdout=True)
     return self.CommonVerifyImage(self, output, percent_required_to_pass)
 
@@ skipping 42 matching lines @@
     Info('Testing for %s' % self.vm_image_path)
 
     self.assertTrue(os.path.exists(self.vm_image_path))
 
   def UpdateImage(self, image_path, stateful_change='old'):
     """Updates VM image with image_path."""
     stateful_change_flag = self.GetStatefulChangeFlag(stateful_change)
     if self.source_image == base_image_path:
       self.source_image = self.vm_image_path
 
-    (code, stdout, stderr) = RunCommandCaptureOutput([
-        '%s/cros_run_vm_update' % self.crosutilsbin,
-        '--update_image_path=%s' % image_path,
-        '--vm_image_path=%s' % self.vm_image_path,
-        '--snapshot',
-        vm_graphics_flag,
-        '--persist',
-        '--kvm_pid=%s' % _KVM_PID_FILE,
-        stateful_change_flag,
-        '--src_image=%s' % self.source_image,
-        ])
-
-    if code != 0:
-      raise UpdateException(code, stdout)
+    cmd = ['%s/cros_run_vm_update' % self.crosutilsbin,
+           '--update_image_path=%s' % image_path,
+           '--vm_image_path=%s' % self.vm_image_path,
+           '--snapshot',
+           vm_graphics_flag,
+           '--persist',
+           '--kvm_pid=%s' % _KVM_PID_FILE,
+           stateful_change_flag,
+           '--src_image=%s' % self.source_image,
+           ]
+    if self.verbose:
+      try:
+        RunCommand(cmd)
+      except Exception, e:
+        raise UpdateException(1, e.message)
+    else:
+      (code, stdout, stderr) = RunCommandCaptureOutput(cmd)
+      if code != 0:
+        raise UpdateException(code, stdout)
 
   def UpdateUsingPayload(self, update_path, stateful_change='old'):
     """Updates a remote image using image_to_live.sh."""
     stateful_change_flag = self.GetStatefulChangeFlag(stateful_change)
     if self.source_image == base_image_path:
       self.source_image = self.vm_image_path
 
-    (code, stdout, stderr) = RunCommandCaptureOutput([
-        '%s/cros_run_vm_update' % self.crosutilsbin,
-        '--payload=%s' % update_path,
-        '--vm_image_path=%s' % self.vm_image_path,
-        '--snapshot',
-        vm_graphics_flag,
-        '--persist',
-        '--kvm_pid=%s' % _KVM_PID_FILE,
-        stateful_change_flag,
-        '--src_image=%s' % self.source_image,
-        ])
+    cmd = ['%s/cros_run_vm_update' % self.crosutilsbin,
+           '--payload=%s' % update_path,
+           '--vm_image_path=%s' % self.vm_image_path,
+           '--snapshot',
+           vm_graphics_flag,
+           '--persist',
+           '--kvm_pid=%s' % _KVM_PID_FILE,
+           stateful_change_flag,
+           '--src_image=%s' % self.source_image,
+           ]
 
-    if code != 0:
-      raise UpdateException(code, stdout)
+    if self.verbose:
+      try:
+        RunCommand(cmd)
+      except Exception, e:
+        raise UpdateException(1, e.message)
+    else:
+      (code, stdout, stderr) = RunCommandCaptureOutput(cmd)
+      if code != 0:
+        raise UpdateException(code, stdout)
 
   def VerifyImage(self, percent_required_to_pass):
     """Runs vm smoke suite to verify image."""
     # image_to_live already verifies lsb-release matching.  This is just
     # for additional steps.
 
     commandWithArgs = ['%s/cros_run_vm_test' % self.crosutilsbin,
                        '--image_path=%s' % self.vm_image_path,
                        '--snapshot',
                        '--persist',
                        '--kvm_pid=%s' % _KVM_PID_FILE,
                        _VERIFY_SUITE,
                        ]
 
     if vm_graphics_flag:
       commandWithArgs.append(vm_graphics_flag)
 
     output = RunCommand(commandWithArgs, error_ok=True, enter_chroot=False,
                         redirect_stdout=True)
     return self.CommonVerifyImage(self, output, percent_required_to_pass)
 
 
 if __name__ == '__main__':
   parser = optparse.OptionParser()
   parser.add_option('-b', '--base_image',
                     help='path to the base image.')
-  parser.add_option('-t', '--target_image',
-                    help='path to the target image.')
   parser.add_option('-r', '--board',
                     help='board for the images.')
-  parser.add_option('-p', '--type', default='vm',
-                    help='type of test to run: [vm, real]. Default: vm.')
-  parser.add_option('-m', '--remote',
-                    help='Remote address for real test.')
-  parser.add_option('--no_graphics', action='store_true',
-                    help='Disable graphics for the vm test.')
   parser.add_option('--no_delta', action='store_false', default=True,
                     dest='delta',
                     help='Disable using delta updates.')
+  parser.add_option('--no_graphics', action='store_true',
+                    help='Disable graphics for the vm test.')
+  parser.add_option('-m', '--remote',
+                    help='Remote address for real test.')
   parser.add_option('-q', '--quick_test', default=False, action='store_true',
                     help='Use a basic test to verify image.')
+  parser.add_option('-t', '--target_image',
+                    help='path to the target image.')
+  parser.add_option('--test_prefix', default='test',
+                    help='Only runs tests with specific prefix i.e. '
+                         'testFullUpdateWipeStateful.')
+  parser.add_option('-p', '--type', default='vm',
+                    help='type of test to run: [vm, real]. Default: vm.')
+  parser.add_option('--verbose', default=False, action='store_true',
+                    help='Print out rather than capture output as much as '
+                         'possible.')
   # Set the usage to include flags.
   parser.set_usage(parser.format_help())
   # Parse existing sys.argv so we can pass rest to unittest.main.
   (options, sys.argv) = parser.parse_args(sys.argv)
 
+  AUTest.verbose = options.verbose
   base_image_path = options.base_image
   target_image_path = options.target_image
   board = options.board
 
   if not base_image_path:
     parser.error('Need path to base image for vm.')
   elif not os.path.exists(base_image_path):
     Die('%s does not exist' % base_image_path)
 
   if not target_image_path:
     parser.error('Need path to target image to update with.')
   elif not os.path.exists(target_image_path):
     Die('%s does not exist' % target_image_path)
 
   if not board:
     parser.error('Need board to convert base image to vm.')
 
   # Communicate flags to tests.
   vm_graphics_flag = ''
   if options.no_graphics: vm_graphics_flag = '--no_graphics'
   if options.quick_test: _VERIFY_SUITE = 'build_RootFilesystemSize'
   AUTest.use_delta_updates = options.delta
 
   # Only run the test harness we care about.
-  if options.type == 'vm':
-    suite = unittest.TestLoader().loadTestsFromTestCase(VirtualAUTest)
-    test_result = unittest.TextTestRunner(verbosity=2).run(suite)
-  elif options.type == 'real':
-    if not options.remote:
-      parser.error('Real tests require a remote test machine.')
-    else:
-      remote = options.remote
+  test_loader = unittest.TestLoader()
+  test_loader.testMethodPrefix = options.test_prefix
 
-    suite = unittest.TestLoader().loadTestsFromTestCase(RealAUTest)
-    test_result = unittest.TextTestRunner(verbosity=2).run(suite)
-  else:
-    parser.error('Could not parse harness type %s.' % options.type)
+  if options.type == 'vm':  test_class = VirtualAUTest
+  elif options.type == 'real': test_class = RealAUTest
+  else: parser.error('Could not parse harness type %s.' % options.type)
+
+  test_suite = test_loader.loadTestsFromTestCase(test_class)
+  test_result = unittest.TextTestRunner(verbosity=2).run(test_suite)
 
   if not test_result.wasSuccessful():
     Die('Test harness was not successful')