Don't include archive diectories that are missing an image.zip when calculating previous image.

This change allows for a previous archive directory to be devoid of an image.zip
without throwing an exception when generating delta payloads. Sometimes we might
be missing an image.zip directory if the state of the last build was bad i.e. it
was aborted midstream -- archive directory is created but image.zip doesn't finish
getting generated.

I've also updated a lot of the logging/docstring to make clear what's happening
in image_extractor.

BUG=chromium-os:32122
TEST=Unittests + pylint + running with full_suite w/ and w/out image.zip in latest.

Change-Id: I5df2451f24721ed34ff8a0a91c76896083ed9a0c
Reviewed-on: https://gerrit.chromium.org/gerrit/29205
Tested-by: Chris Sosa <sosa@chromium.org>
Reviewed-by: David James <davidjames@chromium.org>
Commit-Ready: Chris Sosa <sosa@chromium.org>
diff --git a/generate_test_payloads/cros_generate_test_payloads.py b/generate_test_payloads/cros_generate_test_payloads.py
index 99de657..e968dcd 100755
--- a/generate_test_payloads/cros_generate_test_payloads.py
+++ b/generate_test_payloads/cros_generate_test_payloads.py
@@ -386,12 +386,10 @@
     if latest_image_dir:
       options.base = extractor.UnzipImage(latest_image_dir)
     else:
-      logging.info('Config specified but no previous images found. '
-                   'Using target.')
-      options.base = options.target
+      logging.warning('No previous image.zip found in local archive.')
 
-  elif not options.base:
-    logging.info('Base image not specified. Using target as base image.')
+  if not options.base:
+    logging.info('Using target image as base image.')
     options.base = options.target
 
   if not os.path.isfile(options.base):
diff --git a/lib/image_extractor.py b/lib/image_extractor.py
index 8621202..10c8304 100644
--- a/lib/image_extractor.py
+++ b/lib/image_extractor.py
@@ -35,6 +35,8 @@
         directory may be being populated with the results of this version
         while we're running so we shouldn't use it as the last image archived.
     """
+    logging.info('Searching for previously generated images in %s ... ',
+                 self.archive)
     if os.path.exists(self.archive):
       my_re = re.compile(r'R\d+-(\d+)\.(\d+)\.(\d+).*')
       filelist = []
@@ -42,36 +44,49 @@
       for filename in os.listdir(self.archive):
         lv = distutils.version.LooseVersion(filename)
         if my_re.match(filename):
-          if lv < target_lv:
+          zip_image = os.path.join(self.archive, filename, 'image.zip')
+          if lv < target_lv and os.path.exists(zip_image):
             filelist.append(lv)
           elif not filename.startswith(target_version):
-            logging.error('Version in archive dir is too new: %s' % filename)
+            logging.error('Version in archive dir is too new: %s', filename)
       if filelist:
         return os.path.join(self.archive, str(max(filelist)))
 
+    logging.warn('Could not find a previously generated image on this host.')
     return None
 
   def UnzipImage(self, image_dir):
-    """Unzips the image.zip from the image_dir and returns the image."""
-    # We include the dirname of the image here so that we don't have to
-    # re-unzip the same one each time.
-    local_path = os.path.join(self.SRC_ARCHIVE_DIR,
-                              os.path.basename(image_dir))
-    image_to_return = os.path.abspath(os.path.join(local_path,
-                                                   self.IMAGE_TO_EXTRACT))
-    # We only unzip it if we don't have it.
-    if not os.path.exists(image_to_return):
-      # We don't want to keep test self.SRC_ARCHIVE_DIRs around.
-      if os.path.exists(self.SRC_ARCHIVE_DIR):
-        logging.info('Removing old archive from %s', self.SRC_ARCHIVE_DIR)
-        shutil.rmtree(self.SRC_ARCHIVE_DIR)
+    """Unzips the image.zip from the image_dir and returns the image.
 
-      logging.info('Creating directory %s to store image for testing.',
-                   local_path)
-      os.makedirs(local_path)
+    This method unzips the image under SRC_ARCHIVE_DIR along with its version
+    string. In order to save time, if it is attempting
+    to re-unzip the same image with the same version string, it uses the
+    cached image in SRC_ARCHIVE_DIR. It determines the version string based
+    on the basename of the image_dir.
+
+    Returns: the path to the image.bin file after it has been unzipped.
+    Raises: MissingImageZipException if there is nothing to unzip within
+      the image_dir.
+    """
+    version_string = os.path.basename(image_dir)
+    cached_dir = os.path.join(ImageExtractor.SRC_ARCHIVE_DIR, version_string)
+    cached_image = os.path.abspath(os.path.join(
+        cached_dir, ImageExtractor.IMAGE_TO_EXTRACT))
+    # If we previously unzipped the image, we're done.
+    if os.path.exists(cached_image):
+      logging.info('Re-using image with version %s that we previously '
+                   'unzipped to %s.', version_string, cached_image)
+    else:
+      # Cached image for version not found. Unzipping image from archive.
+      if os.path.exists(ImageExtractor.SRC_ARCHIVE_DIR):
+        logging.info('Removing previously archived images from %s',
+                     ImageExtractor.SRC_ARCHIVE_DIR)
+        shutil.rmtree(ImageExtractor.SRC_ARCHIVE_DIR)
+
+      os.makedirs(cached_dir)
       zip_path = os.path.join(image_dir, 'image.zip')
-      logging.info('Unzipping image from %s', zip_path)
-      chromite_build_lib.RunCommand(['unzip', '-d', local_path, zip_path],
+      logging.info('Unzipping image from %s to %s', zip_path, cached_dir)
+      chromite_build_lib.RunCommand(['unzip', '-d', cached_dir, zip_path],
                                     print_cmd=False)
 
-    return image_to_return
+    return cached_image
diff --git a/lib/image_extractor_unittest.py b/lib/image_extractor_unittest.py
index 6e95f0d..e4d0b3f 100755
--- a/lib/image_extractor_unittest.py
+++ b/lib/image_extractor_unittest.py
@@ -46,6 +46,14 @@
   def tearDown(self):
     shutil.rmtree(self.work_dir)
 
+  @staticmethod
+  def _TouchImageZip(directory):
+    if not os.path.exists(directory):
+      os.makedirs(directory)
+
+    with open(os.path.join(directory, 'image.zip'), 'w') as f:
+      f.write('oogabooga')
+
   def CreateFakeArchiveDir(self, number_of_entries, add_build_number=False):
     """Creates a fake archive dir with specified number of entries."""
     # Create local board directory e.g. /var/archive/x86-generic-full.
@@ -57,6 +65,7 @@
     for i in range(number_of_entries):
       new_dir = os.path.join(self.archive_dir, version_s % i)
       os.makedirs(new_dir)
+      ImageExtractorTest._TouchImageZip(new_dir)
 
   def testGetLatestImageWithNoEntries(self):
     """Should return None if the directory has no entries."""
@@ -111,15 +120,16 @@
   def testUnzipImageArchiveAlready(self):
     """Ensure we create a new archive and delete the old one."""
     old_entry = os.path.join(self.src_archive, 'R16-158.0.0-a1')
-    new_entry = os.path.join(self.src_archive, 'R16-158.0.1-a1')
     os.makedirs(old_entry)
+    new_entry = os.path.join(self.src_archive, 'R16-158.0.1-a1')
+    archived_image_dir = os.path.join(self.archive_dir, 'R16-158.0.1-a1')
+    ImageExtractorTest._TouchImageZip(archived_image_dir)
 
     self.mox.StubOutWithMock(chromite_build_lib, 'RunCommand')
     chromite_build_lib.RunCommand(mox.In('unzip'), print_cmd=False)
 
     self.mox.ReplayAll()
-    self.test_extractor.UnzipImage(os.path.join(self.archive_dir,
-                                                'R16-158.0.1-a1'))
+    self.test_extractor.UnzipImage(archived_image_dir)
     self.mox.VerifyAll()
     self.assertFalse(os.path.exists(old_entry))
     self.assertTrue(os.path.exists(new_entry))
@@ -127,13 +137,14 @@
   def testUnzipImageNoArchive(self):
     """Ensure we create a new archive with none before."""
     new_entry = os.path.join(self.src_archive, 'R16-158.0.1-a1')
+    archived_image_dir = os.path.join(self.archive_dir, 'R16-158.0.1-a1')
+    ImageExtractorTest._TouchImageZip(archived_image_dir)
 
     self.mox.StubOutWithMock(chromite_build_lib, 'RunCommand')
     chromite_build_lib.RunCommand(mox.In('unzip'), print_cmd=False)
 
     self.mox.ReplayAll()
-    self.test_extractor.UnzipImage(os.path.join(self.archive_dir,
-                                                'R16-158.0.1-a1'))
+    self.test_extractor.UnzipImage(archived_image_dir)
     self.mox.VerifyAll()
     self.assertTrue(os.path.exists(new_entry))