Update image_chromeos.py to use 'cros flash'

Check the chroot to see if it contains the 'cros flash' script; if so
use 'cros flash' instead of 'cros_image_to_target.py' for imaging the
DUT.  This is the first of a series of changes to improve the imaging
options in Crosperf.  This script does not change any of the options or
syntax for using image_chromeos.py.

BUG=None
TEST=Tested in chroot with & without cros_flash.py.  Worked correctly in
both cases.

Change-Id: If8704c81aa800cab85379baa514249105fe2b83f
Reviewed-on: https://chrome-internal-review.googlesource.com/153197
Reviewed-by: Han Shen <shenhan@google.com>
Commit-Queue: Caroline Tice <cmtice@google.com>
Tested-by: Caroline Tice <cmtice@google.com>

1 files changed, 25 insertions, 10 deletions

diff --git a/image_chromeos.py b/image_chromeos.py
index 1e061cdd..17588d93 100755
--- a/image_chromeos.py
+++ b/image_chromeos.py
@@ -133,16 +133,31 @@ def DoImage(argv):
     chroot_image = os.path.join(
         "..",
         located_image[len(real_src_dir):].lstrip("/"))
-    cros_image_to_target_args = ["--remote=%s" % options.remote,
-                                 "--board=%s" % board,
-                                 "--from=%s" % os.path.dirname(chroot_image),
-                                 "--image-name=%s" %
-                                 os.path.basename(located_image)]
-
-    command = ("./bin/cros_image_to_target.py %s" %
-               " ".join(cros_image_to_target_args))
-    if options.image_args:
-      command += " %s" % options.image_args
+
+    # Check to see if cros flash is in the chroot or not.
+    cros_flash_path = os.path.join(options.chromeos_root,
+                                   "chromite/cros/commands/cros_flash.py")
+    if os.path.exists(cros_flash_path):
+      # Use 'cros flash'
+      cros_flash_args = ["--board=%s" % board,
+                         "--clobber-stateful",
+                         options.remote,
+                         chroot_image]
+
+      command = ("cros flash %s" % " ".join(cros_flash_args))
+    else:
+      # Use 'cros_image_to_target.py'
+
+      cros_image_to_target_args = ["--remote=%s" % options.remote,
+                                   "--board=%s" % board,
+                                   "--from=%s" % os.path.dirname(chroot_image),
+                                   "--image-name=%s" %
+                                   os.path.basename(located_image)]
+
+      command = ("./bin/cros_image_to_target.py %s" %
+                 " ".join(cros_image_to_target_args))
+      if options.image_args:
+        command += " %s" % options.image_args
 
     # Workaround for crosbug.com/35684.
     os.chmod(misc.GetChromeOSKeyFile(options.chromeos_root), 0600)