cros_image_to_target handles test images, less verbose

bin/cros_image_to_target.py

Index: bin/cros_image_to_target.py
diff --git a/bin/cros_image_to_target.py b/bin/cros_image_to_target.py
index dc6a0f1be5619fbeb7db3e0869c789df269c32e5..cc7d358bbe3fbe95a6a5b95934143fec5f6d9a91 100755
--- a/bin/cros_image_to_target.py
+++ b/bin/cros_image_to_target.py
@@ -1,6 +1,6 @@
 #!/usr/bin/python
 #
-# Copyright (c) 2010 The Chromium OS Authors. All rights reserved.
+# 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.
 
@@ -30,10 +30,11 @@ 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.image.gz'
+STATEFUL_FILENAME = 'stateful.gz'
 
 # How long do we wait for the server to start before launching client
 SERVER_STARTUP_WAIT = 1
@@ -46,8 +47,12 @@ class Command(object):
     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:
@@ -61,8 +66,10 @@ class Command(object):
         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:
@@ -110,9 +117,11 @@ class SSHCommand(Command):
     if not self.ssh_dir:
       self.Setup()
 
+    # allow up to 3 seconds to connect
     return ['-o', 'Compression=no',
             '-o', 'ConnectTimeout=%d' % self.CONNECT_TIMEOUT,
             '-o', 'StrictHostKeyChecking=no',
+            '-o', 'ConnectTimeout=3',

[Paul Stewart, 2010/11/17 21:37:36]

Look 2 lines above this one.  I even used a constant. :-)

[sjg, 2010/11/23 18:49:14]

Sorry! It was hanging at some point and I was looking for why...

             '-o', 'UserKnownHostsFile=%s' % self.known_hosts,
             '-i', self.identity]
 
@@ -139,7 +148,7 @@ class CrosEnv(object):
   REBOOT_START_WAIT = 5
   REBOOT_WAIT_TIME = 60
 
-  def __init__(self, verbose=False):
+  def __init__(self, verbose=0):
     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')):
@@ -147,6 +156,9 @@ class CrosEnv(object):
     self.cmd = Command(self)
     self.verbose = verbose
 
+    # do we have the pv progress tool? (sudo apt-get install pv)
+    self.have_pv = (os.system('pv --help >/dev/null 2>&1') == 0)

[Paul Stewart, 2010/11/17 21:37:36]

I'd much prefer if you used self.cmd.Run() instead

+
   def Error(self, msg):
     print >> sys.stderr, 'ERROR: %s' % msg
 
@@ -156,9 +168,13 @@ class CrosEnv(object):
     sys.exit(1)
 
   def Info(self, msg):
-    if self.verbose:
+    if self.verbose > 0:
       print 'INFO: %s' % msg
 
+  def Debug(self, msg):
+    if self.verbose > 1:
+      print 'DEBUG: %s' % msg
+
   def CrosUtilsPath(self, filename):
     return os.path.join(self.cros_root, filename)
 
@@ -199,6 +215,7 @@ class CrosEnv(object):
       self.Info('Using cached stateful %s' % dst)
       return True
 
+    self.Info('Building stateful')
     cgpt = self.ChrootPath('/usr/bin/cgpt')
     offset = self.cmd.OutputOneLine(cgpt, 'show', '-b', '-i', '1', src)
     size = self.cmd.OutputOneLine(cgpt, 'show', '-s', '-i', '1', src)
@@ -206,9 +223,10 @@ class CrosEnv(object):
       self.Error('Unable to use cgpt to get image geometry')
       return False
 
-    return self.cmd.RunPipe([['dd', 'if=%s' % src, 'bs=512',
-                              'skip=%s' % offset, 'count=%s' % size],
-                             ['gzip', '-c']], outfile=dst)
+    return self.cmd.RunPipe([
+        ['dd', 'if=%s' % src, 'bs=512', 'skip=%s' % offset, 'count=%s' % size],
+        self.have_pv and ['pv', '-s', '%d' % (int (size) * 512)] or ['cat'],
+        ['gzip', '-c']], outfile=dst, hide_stderr=True)
 
   def GetSize(self, filename):
     return os.path.getsize(filename)
@@ -262,10 +280,10 @@ class CrosEnv(object):
     UpdateHandler.SetupUrl('/update', PingUpdateResponse())
     UpdateHandler.SetupUrl('/%s' % UPDATE_FILENAME,
                            FileUpdateResponse(update_file,
-                                              verbose=self.verbose))
+                                              verbose=self.verbose, have_pv=self.have_pv))

[adlr, 2010/11/17 20:43:25]

80 cols here, and a couple lines down as well

     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)
 
@@ -304,6 +322,7 @@ class CrosEnv(object):
   def StartClient(self, port):
     """Ask the client machine to update from our server."""
 
+    self.Info ("Starting client...")

[adlr, 2010/11/17 20:43:25]

i think style says no spaces between method name and (. Also, I think single
quotes are preferred in python

     status = self.GetUpdateStatus()
     if status != 'UPDATE_STATUS_IDLE':
       self.Error('Client update status is not IDLE: %s' % status)
@@ -322,6 +341,7 @@ class CrosEnv(object):
       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,
@@ -334,7 +354,7 @@ class CrosEnv(object):
       self.Error('Client may not have successfully rebooted...')
       return False
 
-    print 'Client update completed successfully!'
+    self.Info ('Client update completed successfully!')

[adlr, 2010/11/17 20:43:25]

remove space before (

     return True
 
 
@@ -342,7 +362,7 @@ 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
 
 
@@ -350,11 +370,12 @@ 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."""
@@ -373,26 +394,33 @@ class FileUpdateResponse(UpdateResponse):
                         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:
-      sent_size = 0
-      sent_percentage = None
-      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
-            print '\rSent %d%%' % sent_percentage,
-            sys.stdout.flush()
+    if send_content:
+      try:
+        fp = False
+        sent_size = 0
+        sent_percentage = None
+
+        # sadly pv gets caught up in the server, not sure why
+        if False: # self.verbose and self.have_pv:
+          fp = subprocess.Popen('pv -s %d >/dev/null' % filesize,
+                shell=True, stdin=subprocess.PIPE, close_fds=True).stdin
+        while True:
+          buf = f.read(self.blocksize)
+          if not buf:
+            break
+          handler.wfile.write(buf)
+          if fp:
+            fp.write(buf)
+          if self.verbose:
+            sent_size += len(buf)
+            percentage = int(100 * sent_size / filesize)
+            if sent_percentage != percentage:
+              sent_percentage = percentage
+              print '\rSent %d%%' % sent_percentage,
+              sys.stdout.flush()
+      finally:
+        if fp:
+          fp.close()

[Paul Stewart, 2010/11/17 21:37:36]

If you can't get it working, please revert this change.

       if self.verbose:
         print '\n'
     f.close()
@@ -556,7 +584,6 @@ def main(argv):
   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')
@@ -565,11 +592,23 @@ def main(argv):
   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
+  maybe_build_test = False

[adlr, 2010/11/17 20:43:25]

it would be nice to have a variable name that more accurately describes its use
(the comment doesn't seem to). would build_test be a better name?

+
+  verbosity = 0
+  if options.verbose:
+      verbosity = 1
+  if options.debug:
+      verbosity = 2
+  cros_env = CrosEnv(verbose=verbosity)
 
   if not options.board:
     options.board = cros_env.GetDefaultBoard()
@@ -584,17 +623,43 @@ def main(argv):
   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
+      maybe_build_test = 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 maybe_build_test:
+	# 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'),

[adlr, 2010/11/17 20:43:25]

a few lines around here (this, 646, 649) are over 80 cols

+                '--', './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:
@@ -610,9 +675,6 @@ def main(argv):
     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, stateful_file)):
     cros_env.Fatal()