chromite: Add ImageTestStage to run tests on built image.

This stage is run in parallel to other stages after BuildImageStage. It
will launch test_image outside the chroot. This stage is forgiving.

Tests live in chromite/cros/tests/image_test.py.

BUG=chromium:382708
BUG=chromium:385355
TEST=unittest
TEST=test_image path/to/image_dir
TEST=test_image path/to/chromium_image.bin
TEST=cbuildbot --local amd64-generic-asan --nobootstrap --noreexec \
       --nouprev --nobuild --noclean --notests
     Make sure that ImageTestStage is launched.

Change-Id: Ia9f1123f9c533ad70a0091ee943d21cddc577ef0
Reviewed-on: https://chromium-review.googlesource.com/203698
Reviewed-by: Aviv Keshet <akeshet@chromium.org>
Reviewed-by: Don Garrett <dgarrett@chromium.org>
Commit-Queue: Nam Nguyen <namnguyen@chromium.org>
Tested-by: Nam Nguyen <namnguyen@chromium.org>
diff --git a/scripts/test_image.py b/scripts/test_image.py
new file mode 100755
index 0000000..093c010
--- /dev/null
+++ b/scripts/test_image.py
@@ -0,0 +1,101 @@
+#!/usr/bin/python
+# Copyright 2014 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.
+
+"""Script to mount a built image and run tests on it."""
+
+import logging
+import os
+import shutil
+import unittest
+
+from chromite.cbuildbot import constants
+from chromite.cros.tests import image_test
+from chromite.lib import commandline
+from chromite.lib import cros_build_lib
+from chromite.lib import osutils
+
+
+def ParseArgs(args):
+  """Return parsed commandline arguments."""
+
+  parser = commandline.ArgumentParser()
+  parser.add_argument('--test_results_root', type='path',
+                      help='Directory to store test results')
+  parser.add_argument('--board', type=str, help='Board (wolf, beaglebone...)')
+  parser.add_argument('image_dir', type='path',
+                      help='Image directory (or file) with mount_image.sh and '
+                           'umount_image.sh')
+  opts = parser.parse_args(args)
+  opts.Freeze()
+  return opts
+
+
+def FindImage(image_path):
+  """Return the path to the image file.
+
+  Args:
+    image_path: A path to the image file, or a directory containing the base
+      image.
+
+  Returns:
+    ImageFileAndMountScripts containing absolute paths to the image,
+      the mount and umount invocation commands
+  """
+
+  if os.path.isdir(image_path):
+    # Assume base image.
+    image_file = os.path.join(image_path, constants.BASE_IMAGE_NAME + '.bin')
+    if not os.path.exists(image_file):
+      raise ValueError('%s does not contain base image' % image_path)
+  elif os.path.isfile(image_path):
+    image_file = image_path
+  else:
+    raise ValueError('%s is neither a directory nor a file' % image_path)
+
+  return image_file
+
+
+def main(args):
+  opts = ParseArgs(args)
+
+  # Build up test suites.
+  loader = unittest.TestLoader()
+  loader.suiteClass = image_test.ImageTestSuite
+  # We use a different prefix here so that unittest DO NOT pick up the
+  # image tests automatically because they depend on a proper environment.
+  loader.testMethodPrefix = 'Test'
+  all_tests = loader.loadTestsFromName('chromite.cros.tests.image_test')
+  forgiving = image_test.ImageTestSuite()
+  non_forgiving = image_test.ImageTestSuite()
+  for suite in all_tests:
+    for test in suite.GetTests():
+      if test.IsForgiving():
+        forgiving.addTest(test)
+      else:
+        non_forgiving.addTest(test)
+
+  # Run them in the image directory.
+  runner = image_test.ImageTestRunner()
+  runner.SetBoard(opts.board)
+  runner.SetResultDir(opts.test_results_root)
+  image_file = FindImage(opts.image_dir)
+  tmp_in_chroot = cros_build_lib.FromChrootPath('/tmp')
+  with osutils.TempDir(base_dir=tmp_in_chroot) as temp_dir:
+    # Copy the image file to a temp dir so that we own it toally, no sharing
+    # in the mount/umount commands.
+    shutil.copy(image_file, temp_dir)
+    image_file = os.path.join(temp_dir, os.path.basename(image_file))
+    with osutils.MountImageContext(image_file, temp_dir):
+      with osutils.ChdirContext(temp_dir):
+        # Run non-forgiving tests first so that exceptions in forgiving tests
+        # do not skip any required tests.
+        logging.info('Running NON-forgiving tests.')
+        result = runner.run(non_forgiving)
+        logging.info('Running forgiving tests.')
+        runner.run(forgiving)
+
+  if result and not result.wasSuccessful():
+    return 1
+  return 0