Crosperf: Remove tar file after uncompressing image

Remove downloaded chromium image.tar.xz file after uncompression.
Change 2 step untar and uncompress to 1. Catch exceptions during
uncompression to avoid having a corrupted bin file if
uncompression is interrupted.

BUG: chromium:663426
TEST: crosperf unit tests pass
Change-Id: I82f674a26bf99cb3491e7db770664c73c004a93b
Reviewed-on: https://chrome-internal-review.googlesource.com/303875
Commit-Ready: Manoj Gupta <manojgupta@google.com>
Tested-by: Manoj Gupta <manojgupta@google.com>
Reviewed-by: Caroline Tice <cmtice@google.com>

2 files changed, 37 insertions, 10 deletions

diff --git a/crosperf/download_images.py b/crosperf/download_images.py
index cabeeb5a..14a91427 100644
--- a/crosperf/download_images.py
+++ b/crosperf/download_images.py
@@ -79,15 +79,32 @@ class ImageDownloader(object):
       return
 
     # Uncompress and untar the downloaded image.
-    command = ('cd /tmp/%s ;unxz chromiumos_test_image.tar.xz; '
-               'tar -xvf chromiumos_test_image.tar' % build_id)
+    command = ('cd /tmp/%s ; tar -Jxf chromiumos_test_image.tar.xz ' % build_id)
     if self.log_level != 'verbose':
       self._logger.LogOutput('CMD: %s' % command)
       print('(Uncompressing and un-tarring may take a couple of minutes...'
             'please be patient.)')
+    try:
+      retval = self._ce.ChrootRunCommand(chromeos_root, command)
+      if retval != 0:
+        raise MissingImage('Cannot uncompress image: %s.' % build_id)
+    except Exception:
+      # Exception in uncompressing file, so cleanup
+      command = ('cd /tmp/%s ; rm -f chromiumos_test_image.bin; ' % build_id)
+      # Remove the partially extracted bin file to avoid issues with future runs
+      _ = self._ce.ChrootRunCommand(chromeos_root, command)
+      # Raise exception again
+      raise
+
+    # Remove uncompressed file
+    command = ('cd /tmp/%s ; rm -f chromiumos_test_image.tar.xz; ' % build_id)
+    if self.log_level != 'verbose':
+      self._logger.LogOutput('CMD: %s' % command)
+      print('(Removing file chromiumos_test_image.tar.xz.)')
+    # try removing file, its ok to have an error, print if encountered
     retval = self._ce.ChrootRunCommand(chromeos_root, command)
     if retval != 0:
-      raise MissingImage('Cannot uncompress image: %s.' % build_id)
+      print('(Warning: Could not remove file chromiumos_test_image.tar.xz .)')
 
   def DownloadSingleAutotestFile(self, chromeos_root, build_id,
                                  package_file_name):
diff --git a/crosperf/download_images_unittest.py b/crosperf/download_images_unittest.py
index fbf2ec79..e2bd7318 100755
--- a/crosperf/download_images_unittest.py
+++ b/crosperf/download_images_unittest.py
@@ -26,6 +26,7 @@ class ImageDownloaderTestcast(unittest.TestCase):
     self.called_download_image = False
     self.called_uncompress_image = False
     self.called_get_build_id = False
+    self.called_download_autotest_files = False
 
   @mock.patch.object(os, 'makedirs')
   @mock.patch.object(os.path, 'exists')
@@ -111,11 +112,18 @@ class ImageDownloaderTestcast(unittest.TestCase):
         '/usr/local/home/chromeos/chroot/tmp/lumpy-release/'
         'R36-5814.0.0/chromiumos_test_image.bin')
 
-    # 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 '
-        'chromiumos_test_image.tar.xz; tar -xvf chromiumos_test_image.tar')
+    # Verify ChrootRunCommand was called twice, with correct arguments.
+    self.assertEqual(mock_cmd_exec.ChrootRunCommand.call_count, 2)
+    call_args_0 = mock_cmd_exec.ChrootRunCommand.call_args_list[0][0]
+    call_args_1 = mock_cmd_exec.ChrootRunCommand.call_args_list[1][0]
+    expected_arg_0 = ('/usr/local/home/chromeos',
+                      'cd /tmp/lumpy-release/R36-5814.0.0 ; '
+                      'tar -Jxf chromiumos_test_image.tar.xz ')
+    expected_arg_1 = ('/usr/local/home/chromeos',
+                      'cd /tmp/lumpy-release/R36-5814.0.0 ; '
+                      'rm -f chromiumos_test_image.bin; ')
+    self.assertEqual(call_args_0, expected_arg_0)
+    self.assertEqual(call_args_1, expected_arg_1)
 
     # Set os.path.exists to always return False and run uncompress.
     mock_path_exists.reset_mock()
@@ -193,7 +201,8 @@ class ImageDownloaderTestcast(unittest.TestCase):
     self.assertTrue(self.called_uncompress_image)
     # Make sure it called DownloadAutotestFiles
     self.assertTrue(self.called_download_autotest_files)
-    # Make sure it returned a autotest path returned from this call
+    # Make sure it returned an image and  autotest path returned from this call
+    self.assertTrue(image_path == 'chromiumos_test_image.bin')
     self.assertTrue(autotest_path == 'autotest')
 
     # Call Run with a non-empty autotest path
@@ -217,7 +226,8 @@ class ImageDownloaderTestcast(unittest.TestCase):
     self.assertRaises(download_images.MissingImage, downloader.Run, test_chroot,
                       test_autotest_path, test_build_id)
 
-    # Verify that UncompressImage and downloadAutotestFiles was not called, since _DownloadImage "failed"
+    # Verify that UncompressImage and downloadAutotestFiles were not called,
+    # since _DownloadImage "failed"
     self.assertTrue(self.called_download_image)
     self.assertFalse(self.called_uncompress_image)
     self.assertFalse(self.called_download_autotest_files)