refactor and add some checks to image downloading code.

The nightly testing is failing in misterious ways while
downloading images for testing. I am adding a check for the
image after "gsutil cp" and, in the process, refactored
the code a little.

BUG=None
TEST=run_tests

Change-Id: I3e84ee9f2ea3b9b66cc96bf723bd1bcb94f26af7
Reviewed-on: https://chrome-internal-review.googlesource.com/248966
Commit-Ready: Luis Lozano <llozano@chromium.org>
Tested-by: Luis Lozano <llozano@chromium.org>
Reviewed-by: Han Shen <shenhan@google.com>

4 files changed, 24 insertions, 36 deletions

diff --git a/crosperf/download_images.py b/crosperf/download_images.py
index 33875208..69c61f91 100644
--- a/crosperf/download_images.py
+++ b/crosperf/download_images.py
@@ -1,4 +1,3 @@
-
 # Copyright (c) 2014, 2015 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.
@@ -52,24 +51,22 @@ class ImageDownloader(object):
 
     # Check to see if the image has already been downloaded.  If not,
     # download the image.
-    status = 0
     if not os.path.exists(image_path):
       command = 'gsutil cp %s /tmp/%s' % (image_name, build_id)
 
       if self.log_level != 'verbose':
         self._logger.LogOutput('CMD: %s' % command)
       status = self._ce.ChrootRunCommand(chromeos_root, command)
+      if status != 0 or not os.path.exists(image_path):
+        raise MissingImage('Cannot download image: %s.' % image_name)
 
-    if status == 0:
-      return image_path
-    else:
-      return None
+    return image_path
 
   def UncompressImage(self, chromeos_root, build_id):
     # Check to see if the file has already been uncompresssed, etc.
     if os.path.exists(os.path.join(chromeos_root, 'chroot/tmp', build_id,
                                    'chromiumos_test_image.bin')):
-      return 0
+      return
 
     # Uncompress and untar the downloaded image.
     command = ('cd /tmp/%s ;unxz chromiumos_test_image.tar.xz; '
@@ -79,7 +76,8 @@ class ImageDownloader(object):
       print('(Uncompressing and un-tarring may take a couple of minutes...'
             'please be patient.)')
     retval = self._ce.ChrootRunCommand(chromeos_root, command)
-    return retval
+    if retval != 0:
+      raise MissingImage('Cannot uncompress image: %s.' % build_id)
 
   def Run(self, chromeos_root, xbuddy_label):
     build_id = self.GetBuildID(chromeos_root, xbuddy_label)
@@ -94,14 +92,11 @@ class ImageDownloader(object):
       status = self._ce.ChrootRunCommand(chromeos_root, cmd)
     if status != 0:
       raise MissingImage('Cannot find official image: %s.' % image_name)
+
     image_path = self.DownloadImage(chromeos_root, build_id, image_name)
-    retval = 0
-    if image_path:
-      retval = self.UncompressImage(chromeos_root, build_id)
-    else:
-      retval = 1
+    self.UncompressImage(chromeos_root, build_id)
 
-    if retval == 0 and self.log_level != 'quiet':
+    if self.log_level != 'quiet':
       self._logger.LogOutput('Using image from %s.' % image_path)
 
-    return retval, image_path
+    return image_path
diff --git a/crosperf/download_images_unittest.py b/crosperf/download_images_unittest.py
index 1a885c61..273a965c 100755
--- a/crosperf/download_images_unittest.py
+++ b/crosperf/download_images_unittest.py
@@ -1,7 +1,6 @@
 #!/usr/bin/python2
 #
 # Copyright 2014 Google Inc.  All Rights Reserved
-
 """Download image unittest."""
 
 from __future__ import print_function
@@ -45,7 +44,8 @@ class ImageDownloaderTestcast(unittest.TestCase):
     # Set os.path.exists to always return False and run downloader
     mock_path_exists.return_value = False
     test_flag.SetTestMode(True)
-    downloader.DownloadImage(test_chroot, test_build_id, image_path)
+    self.assertRaises(download_images.MissingImage, downloader.DownloadImage,
+                      test_chroot, test_build_id, image_path)
 
     # Verify os.path.exists was called twice, with proper arguments.
     self.assertEqual(mock_path_exists.call_count, 2)
@@ -63,8 +63,7 @@ class ImageDownloaderTestcast(unittest.TestCase):
     # Verify we called ChrootRunCommand once, with proper arguments.
     self.assertEqual(mock_cmd_exec.ChrootRunCommand.call_count, 1)
     mock_cmd_exec.ChrootRunCommand.assert_called_with(
-        '/usr/local/home/chromeos',
-        'gsutil cp '
+        '/usr/local/home/chromeos', 'gsutil cp '
         'gs://chromeos-image-archive/lumpy-release/R36-5814.0.0/'
         'chromiumos_test_image.tar.xz'
         ' /tmp/lumpy-release/R36-5814.0.0')
@@ -103,7 +102,8 @@ class ImageDownloaderTestcast(unittest.TestCase):
 
     # Set os.path.exists to always return False and run uncompress.
     mock_path_exists.return_value = False
-    downloader.UncompressImage(test_chroot, test_build_id)
+    self.assertRaises(download_images.MissingImage, downloader.UncompressImage,
+                      test_chroot, test_build_id)
 
     # Verify os.path.exists was called once, with correct arguments.
     self.assertEqual(mock_path_exists.call_count, 1)
@@ -114,8 +114,7 @@ class ImageDownloaderTestcast(unittest.TestCase):
     # Verify ChrootRunCommand was called, with correct arguments.
     self.assertEqual(mock_cmd_exec.ChrootRunCommand.call_count, 1)
     mock_cmd_exec.ChrootRunCommand.assert_called_with(
-        '/usr/local/home/chromeos',
-        'cd /tmp/lumpy-release/R36-5814.0.0 ;unxz '
+        '/usr/local/home/chromeos', 'cd /tmp/lumpy-release/R36-5814.0.0 ;unxz '
         'chromiumos_test_image.tar.xz; tar -xvf chromiumos_test_image.tar')
 
     # Set os.path.exists to always return False and run uncompress.
@@ -159,7 +158,7 @@ class ImageDownloaderTestcast(unittest.TestCase):
       if root or build_id or image_path:
         pass
       self.called_download_image = True
-      return None
+      raise download_images.MissingImage('Could not download image')
 
     def FakeUncompressImage(root, build_id):
       if root or build_id:
@@ -188,7 +187,8 @@ class ImageDownloaderTestcast(unittest.TestCase):
     downloader.DownloadImage = BadDownloadImage
 
     # Call Run again.
-    downloader.Run(test_chroot, test_build_id)
+    self.assertRaises(download_images.MissingImage, downloader.Run, test_chroot,
+                      test_build_id)
 
     # Verify that UncompressImage was not called, since _DownloadImage "failed"
     self.assertTrue(self.called_download_image)
diff --git a/crosperf/settings.py b/crosperf/settings.py
index 3be85e02..0722eb97 100644
--- a/crosperf/settings.py
+++ b/crosperf/settings.py
@@ -74,8 +74,6 @@ class Settings(object):
     else:
       xbuddy_path = '%s/%s' % (prefix, path_str)
     image_downloader = ImageDownloader(l, log_level)
-    retval, image_path = image_downloader.Run(
+    image_path = image_downloader.Run(
         misc.CanonicalizePath(chromeos_root), xbuddy_path)
-    if retval != 0:
-      raise Exception('Unable to find/download xbuddy image.')
     return image_path
diff --git a/crosperf/settings_unittest.py b/crosperf/settings_unittest.py
index e5ccfd46..ee863b80 100755
--- a/crosperf/settings_unittest.py
+++ b/crosperf/settings_unittest.py
@@ -174,7 +174,7 @@ class TestSettings(unittest.TestCase):
   @mock.patch.object(download_images, 'ImageDownloader')
   def test_get_xbuddy_path(self, mock_downloader, mock_run, mock_logger):
 
-    mock_run.return_value = [0, 'fake_xbuddy_translation']
+    mock_run.return_value = 'fake_xbuddy_translation'
     mock_downloader.Run = mock_run
     board = 'lumpy'
     chromeos_root = '/tmp/chromeos'
@@ -194,19 +194,14 @@ class TestSettings(unittest.TestCase):
     self.settings.GetXbuddyPath(official_str, board, chromeos_root, log_level)
     self.assertEqual(mock_run.call_count, 1)
     self.assertEqual(mock_run.call_args_list[0][0],
-                     ('/tmp/chromeos',
-                      'remote/lumpy-release/R34-5417.0.0'))
+                     ('/tmp/chromeos', 'remote/lumpy-release/R34-5417.0.0'))
 
     mock_run.reset_mock()
     self.settings.GetXbuddyPath(xbuddy_str, board, chromeos_root, log_level)
     self.assertEqual(mock_run.call_count, 1)
-    self.assertEqual(mock_run.call_args_list[0][0],
-                     ('/tmp/chromeos',
-                      'remote/lumpy/latest-dev'))
+    self.assertEqual(mock_run.call_args_list[0][0], ('/tmp/chromeos',
+                                                     'remote/lumpy/latest-dev'))
 
-    mock_run.return_value = [1, 'fake_xbuddy_translation']
-    self.assertRaises(Exception, self.settings.GetXbuddyPath, xbuddy_str, board,
-                      chromeos_root, log_level)
     if mock_logger:
       return