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_unittest.py b/scripts/test_image_unittest.py
new file mode 100755
index 0000000..8ecd48c
--- /dev/null
+++ b/scripts/test_image_unittest.py
@@ -0,0 +1,168 @@
+#!/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.
+
+"""Unit tests for the functions in test_image."""
+
+import os
+import tempfile
+import unittest
+
+from chromite.cbuildbot import constants
+from chromite.cros.tests import image_test
+from chromite.lib import cros_build_lib
+from chromite.lib import cros_test_lib
+from chromite.lib import osutils
+from chromite.scripts import test_image
+
+
+class TestImageTest(cros_test_lib.MockTempDirTestCase):
+  """Common class for tests ImageTest.
+
+  This sets up proper directory with test image. The image file is zero-byte.
+  """
+
+  def setUp(self):
+    # create dummy image file
+    self.image_file = os.path.join(self.tempdir,
+                                   constants.BASE_IMAGE_NAME + '.bin')
+    osutils.WriteFile(self.image_file, '')
+    fake_partitions = {
+        1: cros_build_lib.PartitionInfo(1, 0, 0, 0, 'fs', 'STATE', 'flag'),
+        2: cros_build_lib.PartitionInfo(2, 0, 0, 0, 'fs', 'KERN-A', 'flag'),
+        3: cros_build_lib.PartitionInfo(3, 0, 0, 0, 'fs', 'ROOT-A', 'flag'),
+    }
+    self.PatchObject(cros_build_lib, 'GetImageDiskPartitionInfo',
+                     autospec=True, return_value=fake_partitions)
+    self.PatchObject(osutils.MountImageContext, '_Mount', autospec=True)
+    self.PatchObject(osutils.MountImageContext, '_Unmount', autospec=True)
+
+
+class FindImageTest(TestImageTest):
+  """Test FindImage() function."""
+
+  def _testFindOkay(self, image_path):
+    res = test_image.FindImage(image_path)
+    self.assertEqual(
+        res,
+        os.path.join(self.tempdir, constants.BASE_IMAGE_NAME + '.bin')
+    )
+
+  def testFindWithDirectory(self):
+    self._testFindOkay(self.tempdir)
+
+  def testFindWithFile(self):
+    self._testFindOkay(self.image_file)
+
+  def testFindWithInvalid(self):
+    self.assertRaises(ValueError, test_image.FindImage,
+                      os.path.join(self.tempdir, '404'))
+
+  def testFindWithInvalidDirectory(self):
+    os.unlink(self.image_file)
+    self.assertRaises(ValueError, test_image.FindImage,
+                      os.path.join(self.tempdir))
+
+
+class MainTest(TestImageTest):
+  """Test the main invocation of the script."""
+
+  def testChdir(self):
+    """Verify the CWD is in a temp directory."""
+
+    class CwdTest(image_test.NonForgivingImageTestCase):
+      """A dummy test class to verify current working directory."""
+
+      _expected_dir = None
+
+      def SetCwd(self, cwd):
+        self._expected_dir = cwd
+
+      def testExpectedCwd(self):
+        self.assertEqual(self._expected_dir, os.getcwd())
+
+    self.assertNotEqual('/tmp', os.getcwd())
+    os.chdir('/tmp')
+
+    test = CwdTest('testExpectedCwd')
+    suite = image_test.ImageTestSuite()
+    suite.addTest(test)
+    self.PatchObject(unittest.TestLoader, 'loadTestsFromName', autospec=True,
+                     return_value=[suite])
+
+    # Set up the expected directory.
+    expected_dir = os.path.join(self.tempdir, 'my-subdir')
+    os.mkdir(expected_dir)
+    test.SetCwd(expected_dir)
+    self.PatchObject(tempfile, 'mkdtemp', autospec=True,
+                     return_value=expected_dir)
+
+    argv = [self.tempdir]
+    self.assertEqual(0, test_image.main(argv))
+    self.assertEqual('/tmp', os.getcwd())
+
+  def _testForgiveness(self, forgiveness, expected_result):
+
+    class ForgivenessTest(image_test.ImageTestCase):
+      """A dummy test that is sometime forgiving, sometime not.
+
+      Its only test (testFail) always fail.
+      """
+
+      _forgiving = True
+
+      def SetForgiving(self, value):
+        self._forgiving = value
+
+      def IsForgiving(self):
+        return self._forgiving
+
+      def testFail(self):
+        self.fail()
+
+    test = ForgivenessTest('testFail')
+    test.SetForgiving(forgiveness)
+    suite = image_test.ImageTestSuite()
+    suite.addTest(test)
+    self.PatchObject(unittest.TestLoader, 'loadTestsFromName', autospec=True,
+                     return_value=[suite])
+    argv = [self.tempdir]
+    self.assertEqual(expected_result, test_image.main(argv))
+
+  def testForgiving(self):
+    self._testForgiveness(True, 0)
+
+  def testNonForgiving(self):
+    self._testForgiveness(False, 1)
+
+  def testBoardAndDirectory(self):
+    """Verify that "--board", "--test_results_root" are passed to the tests."""
+
+    class AttributeTest(image_test.ForgivingImageTestCase):
+      """Dummy test class to hold board and directory."""
+
+      def testOkay(self):
+        pass
+
+    test = AttributeTest('testOkay')
+    suite = image_test.ImageTestSuite()
+    suite.addTest(test)
+    self.PatchObject(unittest.TestLoader, 'loadTestsFromName', autospec=True,
+                     return_value=[suite])
+    argv = [
+        '--board',
+        'my-board',
+        '--test_results_root',
+        'your-root',
+        self.tempdir
+    ]
+    test_image.main(argv)
+    # pylint: disable=W0212
+    self.assertEqual('my-board', test._board)
+    # pylint: disable=W0212
+    self.assertEqual('your-root', os.path.basename(test._result_dir))
+
+
+if __name__ == '__main__':
+  cros_test_lib.main()