cros_image_to_target handles test images, less verbose

bin/cros_image_to_target.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.
 
 """Create and copy update image to target host.
 
 auto-update and devserver change out from beneath us often enough
 that despite having to duplicate a litte code, it seems that the
@@ skipping 12 matching lines @@
 import sys
 import tempfile
 import time
 import traceback
 
 from xml.dom import minidom
 
 
 # This is the default filename within the image directory to load updates from
 DEFAULT_IMAGE_NAME = 'chromiumos_image.bin'
+DEFAULT_IMAGE_NAME_TEST = 'chromiumos_test_image.bin'
 
 # The filenames we provide to clients to pull updates
 UPDATE_FILENAME = 'update.gz'
 STATEFUL_FILENAME = 'stateful.tgz'
 
 # How long do we wait for the server to start before launching client
 SERVER_STARTUP_WAIT = 1
 
 
 class Command(object):
   """Shell command ease-ups for Python."""
 
   def __init__(self, env):
     self.env = env
 
   def RunPipe(self, pipeline, infile=None, outfile=None,
-              capture=False, oneline=False):
-    """Perform a command pipeline, with optional input/output filenames."""
+              capture=False, oneline=False, hide_stderr=False):
+    """
+    Perform a command pipeline, with optional input/output filenames.
+
+    hide_stderr     Don't allow output of stderr (default False)
+    """
 
     last_pipe = None
     while pipeline:
       cmd = pipeline.pop(0)
       kwargs = {}
       if last_pipe is not None:
         kwargs['stdin'] = last_pipe.stdout
       elif infile:
         kwargs['stdin'] = open(infile, 'rb')
       if pipeline or capture:
         kwargs['stdout'] = subprocess.PIPE
       elif outfile:
         kwargs['stdout'] = open(outfile, 'wb')
+      if hide_stderr:
+        kwargs['stderr'] = open('/dev/null', 'wb')
 
-      self.env.Info('Running: %s' % ' '.join(cmd))
+      self.env.Debug('Running: %s' % ' '.join(cmd))
       last_pipe = subprocess.Popen(cmd, **kwargs)
 
     if capture:
       ret = last_pipe.communicate()[0]
       if not ret:
         return None
       elif oneline:
         return ret.rstrip('\r\n')
       else:
         return ret
@@ skipping 56 matching lines @@
                                   [src, 'root@%s:%s' %
                                    (self.remote, dest)]])
 
 
 class CrosEnv(object):
   """Encapsulates the ChromeOS build system environment functionality."""
 
   REBOOT_START_WAIT = 5
   REBOOT_WAIT_TIME = 60
 
-  def __init__(self, verbose=False):
+  SILENT = 0
+  INFO = 1
+  DEBUG = 2
+
+  def __init__(self, verbose=SILENT):
     self.cros_root = os.path.dirname(os.path.abspath(sys.argv[0]))
     parent = os.path.dirname(self.cros_root)
     if os.path.exists(os.path.join(parent, 'chromeos-common.sh')):
       self.cros_root = parent
     self.cmd = Command(self)
     self.verbose = verbose
 
+    # do we have the pv progress tool? (sudo apt-get install pv)
+    self.have_pv = True
+    try:
+      self.cmd.Output('pv', '--help')
+    except OSError:
+      self.have_pv = False
+
   def Error(self, msg):
     print >> sys.stderr, 'ERROR: %s' % msg
 
   def Fatal(self, msg=None):
     if msg:
       self.Error(msg)
     sys.exit(1)
 
   def Info(self, msg):
-    if self.verbose:
+    if self.verbose >= CrosEnv.INFO:
       print 'INFO: %s' % msg
 
+  def Debug(self, msg):
+    if self.verbose >= CrosEnv.DEBUG:
+      print 'DEBUG: %s' % msg
+
   def CrosUtilsPath(self, filename):
     return os.path.join(self.cros_root, filename)
 
   def ChrootPath(self, filename):
     return self.CrosUtilsPath(os.path.join('..', '..', 'chroot',
                                            filename.strip(os.path.sep)))
 
   def FileOneLine(self, filename):
     return file(filename).read().rstrip('\r\n')
 
@@ skipping 76 matching lines @@
   def CreateServer(self, port, update_file, stateful_file):
     """Start the devserver clone."""
 
     PingUpdateResponse.Setup(self.GetHash(update_file),
                              self.GetSha256(update_file),
                              self.GetSize(update_file))
 
     UpdateHandler.SetupUrl('/update', PingUpdateResponse())
     UpdateHandler.SetupUrl('/%s' % UPDATE_FILENAME,
                            FileUpdateResponse(update_file,
-                                              verbose=self.verbose))
+                                              verbose=self.verbose,
+                                              have_pv=self.have_pv))
     UpdateHandler.SetupUrl('/%s' % STATEFUL_FILENAME,
                            FileUpdateResponse(stateful_file,
-                                              verbose=self.verbose))
+                                              verbose=self.verbose,
+                                              have_pv=self.have_pv))
 
     self.http_server = BaseHTTPServer.HTTPServer(('', port), UpdateHandler)
 
   def StartServer(self):
     self.Info('Starting http server')
     self.http_server.serve_forever()
 
   def GetUpdateStatus(self):
     status = self.ssh_cmd.Output('/usr/bin/update_engine_client', '--status')
     if not status:
@@ skipping 18 matching lines @@
       self.Info('Client has not yet restarted (try %d).  Waiting...' % attempt)
       wait_time = SSHCommand.CONNECT_TIMEOUT - (time.time() - start)
       if wait_time > 0:
         time.sleep(wait_time)
 
     return False
 
   def StartClient(self, port):
     """Ask the client machine to update from our server."""
 
+    self.Info("Starting client...")
     status = self.GetUpdateStatus()
     if status != 'UPDATE_STATUS_IDLE':
       self.Error('Client update status is not IDLE: %s' % status)
       return False
 
     url_base = 'http://localhost:%d' % port
     update_url = '%s/update' % url_base
     fd, update_log = tempfile.mkstemp(prefix='image-to-target-')
     self.Info('Starting update on client.  Client output stored to %s' %
               update_log)
+
+    # this will make the client read the files we have set up
     self.ssh_cmd.Run('/usr/bin/update_engine_client', '--update',
                      '--omaha_url', update_url, remote_tunnel=(port, port),
                      outfile=update_log)
 
     if self.GetUpdateStatus() != 'UPDATE_STATUS_UPDATED_NEED_REBOOT':
       self.Error('Client update failed')
       return False
 
+    self.Info('Update complete - running update script on client')
     self.ssh_cmd.Copy(self.CrosUtilsPath('../platform/dev/stateful_update'),
                       '/tmp')
     if not self.ssh_cmd.Run('/tmp/stateful_update', url_base,
                             remote_tunnel=(port, port)):
       self.Error('Client stateful update failed')
       return False
 
     self.Info('Rebooting client')
     if not self.ClientReboot():
       self.Error('Client may not have successfully rebooted...')
       return False
 
-    print 'Client update completed successfully!'
+    self.Info('Client update completed successfully!')
     return True
 
 
 class UpdateResponse(object):
   """Default response is the 404 error response."""
 
   def Reply(self, handler, send_content=True, post_data=None):
-    handler.send_Error(404, 'File not found')
+    handler.send_error(404, 'File not found')
     return None
 
 
 class FileUpdateResponse(UpdateResponse):
   """Respond by sending the contents of a file."""
 
   def __init__(self, filename, content_type='application/octet-stream',
-               verbose=False, blocksize=16*1024):
+               verbose=False, blocksize=16*1024, have_pv=False):
     self.filename = filename
     self.content_type = content_type
     self.verbose = verbose
     self.blocksize = blocksize
+    self.have_pv = have_pv
 
   def Reply(self, handler, send_content=True, post_data=None):
     """Return file contents to the client.  Optionally display progress."""
 
     try:
       f = open(self.filename, 'rb')
     except IOError:
       return UpdateResponse.Reply(self, handler)
 
     handler.send_response(200)
     handler.send_header('Content-type', self.content_type)
     filestat = os.fstat(f.fileno())
     filesize = filestat[6]
     handler.send_header('Content-Length', str(filesize))
     handler.send_header('Last-Modified',
                         handler.date_time_string(filestat.st_mtime))
     handler.end_headers()
 
-    if not send_content:
-      return
-
-    if filesize <= self.blocksize:
-      handler.wfile.write(f.read())
-    else:
+    if send_content:
       sent_size = 0
       sent_percentage = None
+
+      #TODO(sjg): this should use pv also
       while True:
         buf = f.read(self.blocksize)
         if not buf:
           break
         handler.wfile.write(buf)
         if self.verbose:
           sent_size += len(buf)
           percentage = int(100 * sent_size / filesize)
           if sent_percentage != percentage:
             sent_percentage = percentage
@@ skipping 155 matching lines @@
   usage = 'usage: %prog [options]'
   parser = optparse.OptionParser(usage=usage)
   parser.add_option('--board', dest='board', default=None,
                     help='Board platform type')
   parser.add_option('--force-mismatch', dest='force_mismatch', default=False,
                     action='store_true',
                     help='Upgrade even if client arch does not match')
   parser.add_option('--from', dest='src', default=None,
                     help='Source image to install')
   parser.add_option('--image-name', dest='image_name',
-                    default=DEFAULT_IMAGE_NAME,
                     help='Filename within image directory to load')
   parser.add_option('--port', dest='port', default=8081, type='int',
                     help='TCP port to serve from and tunnel through')
   parser.add_option('--remote', dest='remote', default=None,
                     help='Remote device-under-test IP address')
   parser.add_option('--server-only', dest='server_only', default=False,
                     action='store_true', help='Do not start client')
   parser.add_option('--verbose', dest='verbose', default=False,
+                    action='store_true', help='Display progress')
+  parser.add_option('--debug', dest='debug', default=False,
                     action='store_true', help='Display running commands')
+  parser.add_option('--test', dest='test', default=False,
+                    action='store_true', help='Select test image')
 
   (options, args) = parser.parse_args(argv)
 
-  cros_env = CrosEnv(verbose=options.verbose)
+  # we can build the test image if it doesn't exist, so remember if we want to
+  build_test_image = False
+
+  verbosity = CrosEnv.SILENT
+  if options.verbose:
+      verbosity = CrosEnv.INFO
+  if options.debug:
+      verbosity = CrosEnv.DEBUG
+  cros_env = CrosEnv(verbose=verbosity)
 
   if not options.board:
     options.board = cros_env.GetDefaultBoard()
 
   if not options.src:
     options.src = cros_env.GetLatestImage(options.board)
     if options.src is None:
       parser.error('No --from argument given and no default image found')
 
   cros_env.Info('Performing update from %s' % options.src)
 
   if not os.path.exists(options.src):
     parser.error('Path %s does not exist' % options.src)
 
+  if not options.image_name:
+    # auto-select the correct image
+    if options.test:
+      options.image_name = DEFAULT_IMAGE_NAME_TEST
+
+      # we will build the test image if not found
+      build_test_image = True
+    else:
+      options.image_name = DEFAULT_IMAGE_NAME
+
   if os.path.isdir(options.src):
     image_directory = options.src
     image_file = os.path.join(options.src, options.image_name)
 
     if not os.path.exists(image_file):
+      if build_test_image:
+        # we want a test image but it doesn't exist
+        # try to build it if we can
+        cros_env.Info('Creating test image')
+        test_output = cros_env.cmd.Output(
+                cros_env.CrosUtilsPath('enter_chroot.sh'),
+                '--', './mod_image_for_test.sh',
+                '--board=%s' % options.board, '-y')
+        if not os.path.exists(image_file):
+          print test_output
+          cros_env.Fatal('Failed to create test image - please run '
+             './mod_image_for_test.sh manually inside the chroot')
       parser.error('Image file %s does not exist' % image_file)
   else:
     image_file = options.src
     image_directory = os.path.dirname(options.src)
 
+  update_file = os.path.join(image_directory, UPDATE_FILENAME)
+  stateful_file = os.path.join(image_directory, STATEFUL_FILENAME)
+
+  cros_env.Debug("Image file %s" % image_file)
+  cros_env.Debug("Update file %s" % update_file)
+  cros_env.Debug("Stateful file %s" % stateful_file)
+
   if options.remote:
+    cros_env.Info('Contacting client %s' % options.remote)
     cros_env.SetRemote(options.remote)
     rel = cros_env.GetRemoteRelease()
     if not rel:
       cros_env.Fatal('Could not retrieve remote lsb-release')
     board = rel.get('CHROMEOS_RELEASE_BOARD', '(None)')
     if board != options.board and not options.force_mismatch:
       cros_env.Error('Board %s does not match expected %s' %
                      (board, options.board))
       cros_env.Error('(Use --force-mismatch option to override this)')
       cros_env.Fatal()
 
   elif not options.server_only:
     parser.error('Either --server-only must be specified or '
                  '--remote=<client> needs to be given')
 
-  update_file = os.path.join(image_directory, UPDATE_FILENAME)
-  stateful_file = os.path.join(image_directory, STATEFUL_FILENAME)
-
   if (not cros_env.GenerateUpdatePayload(image_file, update_file) or
       not cros_env.BuildStateful(image_file, image_directory, stateful_file)):
     cros_env.Fatal()
 
   cros_env.CreateServer(options.port, update_file, stateful_file)
 
   exit_status = 1
   if options.server_only:
     child = None
   else:
@@ skipping 31 matching lines @@
 
   if child:
     os.kill(child, 15)
 
   cros_env.Info('Server exiting with status %d' % exit_status)
   sys.exit(exit_status)
 
 
 if __name__ == '__main__':
   main(sys.argv)