Pre-generate updates for test harness in parallel.

bin/cros_au_test_harness.py

 #!/usr/bin/python
 
 # Copyright (c) 2011 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.
 
 """This module runs a suite of Auto Update tests.
 
   The tests can be run on either a virtual machine or actual device depending
   on parameters given.  Specific tests can be run by invoking --test_prefix.
   Verbose is useful for many of the tests if you want to see individual commands
   being run during the update process.
 """
 
 import optparse
 import os
 import re
 import sys
-import thread
+import threading
 import time
 import unittest
 import urllib
 
 sys.path.append(os.path.join(os.path.dirname(__file__), '../lib'))
 from cros_build_lib import Die
 from cros_build_lib import Info
 from cros_build_lib import ReinterpretPathForChroot
 from cros_build_lib import RunCommand
 from cros_build_lib import RunCommandCaptureOutput
@@ skipping 97 matching lines @@
       raise
 
   def AttemptUpdateWithPayloadExpectedFailure(self, payload, expected_msg):
     """Attempt a payload update, expect it to fail with expected log"""
     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
+      else:
+        Warning("Didn't find '%s' in:" % expected_msg)
+        Warning(err.stdout)
 
-    Warning("Didn't find '%s' in:" % expected_msg)
-    Warning(err.stdout)
     self.fail('We managed to update when failure was expected')
 
   def AttemptUpdateWithFilter(self, filter, proxy_port=8081):
     """Update through a proxy, with a specified filter, and expect success."""
 
     self.PrepareBase(self.target_image_path)
 
     # The devserver runs at port 8080 by default. We assume that here, and
     # start our proxy at a different. We then tell our update tools to
     # have the client connect to our proxy_port instead of 8080.
     proxy = cros_test_proxy.CrosTestProxy(port_in=proxy_port,
                                           address_out='127.0.0.1',
                                           port_out=8080,
                                           filter=filter)
     proxy.serve_forever_in_thread()
 
     # This update is expected to fail...
     try:
-      self.PerformUpdate(self.target_image_path, self.base_image_path,
+      self.PerformUpdate(self.target_image_path, self.target_image_path,
                          proxy_port=proxy_port)
     finally:
       proxy.shutdown()
 
   # -------- Functions that subclasses should override ---------
 
   @classmethod
   def ProcessOptions(cls, parser, options):
     """Processes options.
 
@@ skipping 72 matching lines @@
     rather than wiping it.
     """
     # Just make sure some tests pass on original image.  Some old images
     # don't pass many tests.
     self.PrepareBase(self.base_image_path)
     # TODO(sosa): move to 100% once we start testing using the autotest paired
     # with the dev channel.
     percent_passed = self.VerifyImage(10)
 
     # Update to - all tests should pass on new image.
-    Info('Updating from base image on vm to target image.')
     self.PerformUpdate(self.target_image_path, self.base_image_path)
     self.VerifyImage(100)
 
     # Update from - same percentage should pass that originally passed.
-    Info('Updating from updated image on vm back to base image.')
     self.PerformUpdate(self.base_image_path, self.target_image_path)
     self.VerifyImage(percent_passed)
 
   def testUpdateWipeStateful(self):
     """Tests if we can update after cleaning the stateful partition.
 
     This test checks that we can update successfully after wiping the
     stateful partition.
     """
     # Just make sure some tests pass on original image.  Some old images
     # don't pass many tests.
     self.PrepareBase(self.base_image_path)
     # TODO(sosa): move to 100% once we start testing using the autotest paired
     # with the dev channel.
     percent_passed = self.VerifyImage(10)
 
     # Update to - all tests should pass on new image.
-    Info('Updating from base image on vm to target image and wiping stateful.')
     self.PerformUpdate(self.target_image_path, self.base_image_path, 'clean')
     self.VerifyImage(100)
 
     # Update from - same percentage should pass that originally passed.
-    Info('Updating from updated image back to base image and wiping stateful.')
     self.PerformUpdate(self.base_image_path, self.target_image_path, 'clean')
     self.VerifyImage(percent_passed)
 
   # TODO(sosa): Get test to work with verbose.
   def NotestPartialUpdate(self):
     """Tests what happens if we attempt to update with a truncated payload."""
     # Preload with the version we are trying to test.
     self.PrepareBase(self.target_image_path)
 
     # Image can be updated at:
@@ skipping 211 matching lines @@
 
     # Communicate flags to tests.
     cls.graphics_flag = ''
     if options.no_graphics: cls.graphics_flag = '--no_graphics'
 
     if not cls.board:
       parser.error('Need board to convert base image to vm.')
 
   def PrepareBase(self, image_path):
     """Creates an update-able VM based on base image."""
+    # Needed for VM delta updates.  We need to use the qemu image rather
+    # than the base image on a first update.  By tracking the first_update
+    # we can set src_image to the qemu form of the base image when
+    # performing generating the delta payload.
+    self._first_update = True
     self.vm_image_path = '%s/chromiumos_qemu_image.bin' % os.path.dirname(
         image_path)
-
-    Info('Creating: %s' % self.vm_image_path)
-
     if not os.path.exists(self.vm_image_path):
-      Info('Qemu image %s not found, creating one.' % self.vm_image_path)
+      Info('Creating %s' % vm_image_path)
       RunCommand(['%s/image_to_vm.sh' % self.crosutils,
                   '--full',
                   '--from=%s' % ReinterpretPathForChroot(
                       os.path.dirname(image_path)),
                   '--vdisk_size=%s' % self._FULL_VDISK_SIZE,
                   '--statefulfs_size=%s' % self._FULL_STATEFULFS_SIZE,
                   '--board=%s' % self.board,
                   '--test_image'], enter_chroot=True)
-    else:
-      Info('Using existing VM image %s' % self.vm_image_path)
 
-    Info('Testing for %s' % self.vm_image_path)
+    Info('Using %s as base' % self.vm_image_path)
     self.assertTrue(os.path.exists(self.vm_image_path))
 
   def _UpdateImage(self, image_path, src_image_path='', stateful_change='old',
                    proxy_port=None):
     """Updates VM image with image_path."""
     stateful_change_flag = self.GetStatefulChangeFlag(stateful_change)
-    if src_image_path == self.base_image_path:
+    if src_image_path and self._first_update:
       src_image_path = self.vm_image_path
+      self._first_update = False
 
     cmd = ['%s/cros_run_vm_update' % self.crosutilsbin,
            '--update_image_path=%s' % image_path,
            '--vm_image_path=%s' % self.vm_image_path,
            '--snapshot',
            self.graphics_flag,
            '--persist',
            '--kvm_pid=%s' % self._KVM_PID_FILE,
            stateful_change_flag,
            '--src_image=%s' % src_image_path,
@@ skipping 53 matching lines @@
                        ]
 
     if self.graphics_flag:
       commandWithArgs.append(self.graphics_flag)
 
     output = RunCommand(commandWithArgs, error_ok=True, enter_chroot=False,
                         redirect_stdout=True)
     return self.AssertEnoughTestsPassed(self, output, percent_required_to_pass)
 
 
+class GenerateVirtualAUDeltasTest(VirtualAUTest):
+  """Class the overrides VirtualAUTest and stores deltas we will generate."""
+  delta_list = {}
+
+  def setUp(self):
+    AUTest.setUp(self)
+
+  def _UpdateImage(self, image_path, src_image_path='', stateful_change='old',
+                   proxy_port=None):
+    if src_image_path and self._first_update:
+      src_image_path = self.vm_image_path
+      self._first_update = False
+
+    image_path = ReinterpretPathForChroot(image_path)
+    if src_image_path:
+      src_image_path = ReinterpretPathForChroot(src_image_path)
+    if not self.delta_list.has_key(image_path):
+      self.delta_list[image_path] = set([src_image_path])
+    else:
+      self.delta_list[image_path].add(src_image_path)
+
+  def AttemptUpdateWithPayloadExpectedFailure(self, payload, expected_msg):
+    pass
+
+  def VerifyImage(self, percent_required_to_pass):
+    pass
+
+
+class ParallelJob(threading.Thread):
+  """Small wrapper for threading.Thread that releases a semaphore on exit."""
+  def __init__(self, semaphore, target, args):
+    threading.Thread.__init__(self, target=target, args=args)
+    self._target = target
+    self._args = args
+    self._semaphore = semaphore
+    self._output = None
+    self._completed = False
+
+  def run(self):
+    try:
+      threading.Thread.run(self)
+    finally:
+      self._Cleanup()
+      self._completed = True
+
+  def GetOutput(self):
+    assert self._completed, 'GetOutput called before thread was run.'
+    return self._output
+
+  def _Cleanup(self):
+    self._semaphore.release()
+
+  def __str__(self):
+    return '%s(%s)' % (self._target, self._args)
+
+
+def _RunParallelJobs(number_of_sumultaneous_jobs, jobs, jobs_args):
+  """Runs set number of specified jobs in parallel.
+
+  Args:
+    number_of_simultaneous_jobs:  Max number of threads to be run in parallel.
+    jobs:  Array of methods to run.
+    jobs_args:  Array of args associated with method calls.
+  Returns:
+    Returns an array of results corresponding to each thread.
+  """
+  def _TwoTupleize(x, y):
+    return (x, y)
+
+  threads = []
+  job_pool_semaphore = threading.Semaphore(number_of_sumultaneous_jobs)
+  assert len(jobs) == len(jobs_args), 'Length of args array is wrong.'
+
+  # Create the parallel jobs.
+  for job, args in map(_TwoTupleize, jobs, jobs_args):
+    thread = ParallelJob(job_pool_semaphore, target=job, args=args)
+    threads.append(thread)
+
+  # We use a semaphore to ensure we don't run more jobs that required.
+  # After each thread finishes, it releases (increments semaphore).
+  # Acquire blocks of num jobs reached and continues when a thread finishes.
+  for next_thread in threads:
+    job_pool_semaphore.acquire(blocking=True)
+    Info('Starting %s' % next_thread)
+    next_thread.start()
+
+  # Wait on the rest of the threads to finish.
+  for thread in threads:
+    thread.join()
+
+  return [thread.GetOutput() for thread in threads]
+
+
+def _PregenerateUpdates(parser, options):
+  """Determines all deltas that will be generated and generates them.
+
+  This method effectively pre-generates the dev server cache for all tests.
+
+  Args:
+    parser: parser from main.
+    options: options from parsed parser.
+  Returns:
+    Array of output from generating updates.
+  """
+  def _GenerateVMUpdate(target, src):
+    """Generates an update using the devserver."""
+    RunCommand(['./start_devserver',
+                '--pregenerate_update',
+                '--exit',
+                '--image=%s' % target,
+                '--src_image=%s' % src,
+                '--for_vm'
+               ])
+
+  # Get the list of deltas by mocking out update method in test class.
+  GenerateVirtualAUDeltasTest.ProcessOptions(parser, options)
+  test_loader = unittest.TestLoader()
+  test_loader.testMethodPrefix = options.test_prefix
+  test_suite = test_loader.loadTestsFromTestCase(GenerateVirtualAUDeltasTest)
+  test_result = unittest.TextTestRunner(verbosity=0).run(test_suite)
+
+  Info('The following delta updates are required.')
+  jobs = []
+  args = []
+  for target, srcs in GenerateVirtualAUDeltasTest.delta_list.items():
+    for src in srcs:
+      if src:
+        print >> sys.stderr, 'DELTA AU %s -> %s' % (src, target)
+      else:
+        print >> sys.stderr, 'FULL AU %s' % target
+
+      jobs.append(_GenerateVMUpdate)
+      args.append((target, src))
+
+  results = _RunParallelJobs(options.jobs, jobs, args)
+  return results
+
+
 def main():
   parser = optparse.OptionParser()
   parser.add_option('-b', '--base_image',
                     help='path to the base image.')
   parser.add_option('-r', '--board',
                     help='board for the images.')
   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('-j', '--jobs', default=8, type=int,
+                     help='Number of simultaneous jobs')
+  parser.add_option('-q', '--quick_test', default=False, action='store_true',
+                    help='Use a basic test to verify image.')
   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=True, action='store_true',
                     help='Print out rather than capture output as much as '
                          'possible.')
   (options, leftover_args) = parser.parse_args()
 
   if leftover_args:
     parser.error('Found extra options we do not support: %s' % leftover_args)
 
   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)
 
+  # TODO(sosa): Caching doesn't really make sense on non-vm images (yet).
+  if options.type == 'vm':
+    _PregenerateUpdates(parser, options)
+
+  # Run the test suite.
   test_class.ProcessOptions(parser, options)
-
   test_loader = unittest.TestLoader()
   test_loader.testMethodPrefix = options.test_prefix
   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')
+    Die('Test harness was not successful.')
 
 
 if __name__ == '__main__':
   main()