cros_run_unit_tests: Convert to python.

This replaces the old bash script installed by crostestutils.

BUG=chromium:486263
TEST=unittests.
TEST=trybot run on gizmo.
TEST=`cros_run_unit_tests --pretend --board=gizmo` test the same
  packages with the old and new script.

Change-Id: Idfa3501c8639b41d97165ff80bb06522ed734c3a
Reviewed-on: https://chromium-review.googlesource.com/270148
Trybot-Ready: Bertrand Simonnet <bsimonnet@chromium.org>
Tested-by: Bertrand Simonnet <bsimonnet@chromium.org>
Reviewed-by: Bertrand Simonnet <bsimonnet@chromium.org>
Commit-Queue: Bertrand Simonnet <bsimonnet@chromium.org>

8 files changed, 236 insertions, 3 deletions

diff --git a/bin/cros_run_unit_tests b/bin/cros_run_unit_tests
new file mode 120000
index 000000000..72196ceea
--- /dev/null
+++ b/bin/cros_run_unit_tests
@@ -0,0 +1 @@
+../scripts/wrapper.py
\ No newline at end of file
diff --git a/cbuildbot/commands.py b/cbuildbot/commands.py
index 0db02c8d8..6f8fd69d5 100644
--- a/cbuildbot/commands.py
+++ b/cbuildbot/commands.py
@@ -604,7 +604,8 @@ def RunUnitTests(buildroot, board, full, blacklist=None, extra_env=None):
   if blacklist:
     cmd += ['--blacklist_packages=%s' % ' '.join(blacklist)]
 
-  RunBuildScript(buildroot, cmd, enter_chroot=True, extra_env=extra_env or {})
+  RunBuildScript(buildroot, cmd, chromite_cmd=True, enter_chroot=True,
+                 extra_env=extra_env or {})
 
 
 def RunTestSuite(buildroot, board, image_dir, results_dir, test_type,
diff --git a/lib/chroot_util.py b/lib/chroot_util.py
index 3dde021af..34262aac7 100644
--- a/lib/chroot_util.py
+++ b/lib/chroot_util.py
@@ -174,3 +174,28 @@ def InitializeSysroots(blueprint):
   sysroot.WriteConfig(sysroot.GenerateBrickConfig(brick_stack, bsp))
   sysroot.GeneratePortageConfig()
   sysroot.UpdateToolchain()
+
+
+def RunUnittests(sysroot, packages, extra_env=None):
+  """Runs the unit tests for |packages|.
+
+  Args:
+    sysroot: Path to the sysroot to build the tests in.
+    packages: List of packages to test.
+    extra_env: Python dictionary containing the extra environment variable to
+      pass to the build command.
+
+  Raises:
+    RunCommandError if the unit tests failed.
+  """
+  env = extra_env.copy() if extra_env else {}
+  env.update({
+      'FEATURES': 'test',
+      'PKGDIR': os.path.join(sysroot, 'test-packages'),
+  })
+
+  command = ([os.path.join(constants.CHROMITE_BIN_DIR, 'parallel_emerge'),
+              '--sysroot=%s' % sysroot, '--nodeps', '--buildpkgonly'] +
+             list(packages))
+
+  cros_build_lib.SudoRunCommand(command, extra_env=env)
diff --git a/lib/portage_util.py b/lib/portage_util.py
index d3a98c724..6273e11bd 100644
--- a/lib/portage_util.py
+++ b/lib/portage_util.py
@@ -23,6 +23,7 @@ from chromite.lib import cros_logging as logging
 from chromite.lib import gerrit
 from chromite.lib import git
 from chromite.lib import osutils
+from chromite.lib import parallel
 
 
 # The parsed output of running `ebuild <ebuild path> info`.
@@ -55,6 +56,10 @@ _blank_or_eapi_re = re.compile(r'^\s*(?:#|EAPI=|$)')
 WORKON_EBUILD_VERSION = '9999'
 WORKON_EBUILD_SUFFIX = '-%s.ebuild' % WORKON_EBUILD_VERSION
 
+UNITTEST_PACKAGE_BLACKLIST = set((
+    'sys-devel/binutils',
+))
+
 
 class MissingOverlayException(Exception):
   """This exception indicates that a needed overlay is missing."""
@@ -459,6 +464,7 @@ class EBuild(object):
     self.is_workon = False
     self.is_stable = False
     self.is_blacklisted = False
+    self.has_test = False
     self._ReadEBuild(path)
 
   @staticmethod
@@ -473,6 +479,7 @@ class EBuild(object):
     is_workon = False
     is_stable = False
     is_blacklisted = False
+    has_test = False
     for line in fileinput.input(ebuild_path):
       if line.startswith('inherit ') and 'cros-workon' in line:
         is_workon = True
@@ -482,15 +489,19 @@ class EBuild(object):
             is_stable = True
       elif line.startswith('CROS_WORKON_BLACKLIST='):
         is_blacklisted = True
+      elif (line.startswith('src_test()') or
+            line.startswith('platform_pkg_test()')):
+        has_test = True
     fileinput.close()
-    return is_workon, is_stable, is_blacklisted
+    return is_workon, is_stable, is_blacklisted, has_test
 
   def _ReadEBuild(self, path):
     """Determine the settings of `is_workon`, `is_stable` and is_blacklisted
 
     These are determined using the static Classify function.
     """
-    self.is_workon, self.is_stable, self.is_blacklisted = EBuild.Classify(path)
+    (self.is_workon, self.is_stable,
+     self.is_blacklisted, self.has_test) = EBuild.Classify(path)
 
   @staticmethod
   def GetCrosWorkonVars(ebuild_path, pkg_name):
@@ -1597,3 +1608,36 @@ def CleanOutdatedBinaryPackages(sysroot):
   """Cleans outdated binary packages from |sysroot|."""
   return cros_build_lib.RunCommand(
       [cros_build_lib.GetSysrootToolPath(sysroot, 'eclean'), '-d', 'packages'])
+
+
+def _CheckHasTest(cp, sysroot):
+  """Checks if the ebuild for |cp| has tests.
+
+  Args:
+    cp: A portage package in the form category/package_name.
+    sysroot: Path to the sysroot.
+
+  Returns:
+    |cp| if the ebuild for |cp| defines a test stanza, None otherwise.
+  """
+  ebuild = EBuild(FindEbuildForPackage(cp, sysroot))
+  return cp if ebuild.has_test else None
+
+
+def PackagesWithTest(sysroot, packages):
+  """Returns the subset of |packages| that have unit tests.
+
+  Args:
+    sysroot: Path to the sysroot.
+    packages: List of packages to filter.
+
+  Returns:
+    The subset of |packages| that defines unit tests.
+  """
+  inputs = [(cp, sysroot) for cp in packages]
+  pkg_with_test = set(parallel.RunTasksInProcessPool(_CheckHasTest, inputs))
+
+  # CheckHasTest will return None for packages that do not have tests. We can
+  # discard that value.
+  pkg_with_test.discard(None)
+  return pkg_with_test
diff --git a/lib/portage_util_unittest.py b/lib/portage_util_unittest.py
index 259d61ea6..3470f5d8a 100644
--- a/lib/portage_util_unittest.py
+++ b/lib/portage_util_unittest.py
@@ -45,6 +45,26 @@ class _DummyCommandResult(object):
 
 class EBuildTest(cros_test_lib.MoxTestCase):
   """Ebuild related tests."""
+  _MULTILINE_WITH_TEST = """
+hello
+src_test() {
+}"""
+
+  _MULTILINE_NO_TEST = """
+hello
+src_compile() {
+}"""
+
+  _MULTILINE_COMMENTED = """
+#src_test() {
+# notactive
+# }"""
+
+  _MULTILINE_PLATFORM = """
+platform_pkg_test() {
+}"""
+
+  _SINGLE_LINE_TEST = 'src_test() { echo "foo" }'
 
   def _makeFakeEbuild(self, fake_ebuild_path, fake_ebuild_content=''):
     self.mox.StubOutWithMock(fileinput, 'input')
@@ -142,6 +162,21 @@ class EBuildTest(cros_test_lib.MoxTestCase):
         fake_ebuild_path, fake_ebuild_content=['CROS_WORKON_BLACKLIST="1"\n'])
     self.assertEquals(fake_ebuild.is_blacklisted, True)
 
+  def testHasTest(self):
+    """Tests that we detect test stanzas correctly."""
+    def run_case(content, expected):
+      with osutils.TempDir() as temp:
+        ebuild = os.path.join(temp, 'overlay', 'app-misc',
+                              'foo-0.0.1-r1.ebuild')
+        osutils.WriteFile(ebuild, content, makedirs=True)
+        self.assertEqual(expected, portage_util.EBuild(ebuild).has_test)
+
+    run_case(self._MULTILINE_WITH_TEST, True)
+    run_case(self._MULTILINE_NO_TEST, False)
+    run_case(self._MULTILINE_COMMENTED, False)
+    run_case(self._MULTILINE_PLATFORM, True)
+    run_case(self._SINGLE_LINE_TEST, True)
+
 
 class ProjectAndPathTest(cros_test_lib.MoxTempDirTestCase):
   """Project and Path related tests."""
diff --git a/lib/workon_helper.py b/lib/workon_helper.py
index 34e9f030d..176bef8e6 100644
--- a/lib/workon_helper.py
+++ b/lib/workon_helper.py
@@ -691,3 +691,11 @@ class WorkonHelper(object):
       atoms = self._GetCanonicalAtoms(packages)
     for atom in atoms:
       self.RunCommandInAtomSourceDirectory(atom, command)
+
+  def InstalledWorkonAtoms(self):
+    """Returns the set of installed cros_workon packages."""
+    installed_cp = set()
+    for pkg in portage_util.PortageDB(self._sysroot).InstalledPackages():
+      installed_cp.add('%s/%s' % (pkg.category, pkg.package))
+
+    return set(a for a in self.ListAtoms(use_all=True) if a in installed_cp)
diff --git a/lib/workon_helper_unittest.py b/lib/workon_helper_unittest.py
index 99cf7a85c..8b3271e06 100644
--- a/lib/workon_helper_unittest.py
+++ b/lib/workon_helper_unittest.py
@@ -6,6 +6,7 @@
 
 from __future__ import print_function
 
+import collections
 import os
 
 from chromite.lib import cros_test_lib
@@ -32,6 +33,10 @@ BOARD_OVERLAY_DIR = 'overlay-' + BOARD
 HOST_OVERLAY_DIR = 'overlay-host'
 
 
+InstalledPackageMock = collections.namedtuple('InstalledPackage',
+                                              ('category', 'package'))
+
+
 class WorkonHelperTest(cros_test_lib.MockTempDirTestCase):
   """Tests for chromite.lib.workon_helper."""
 
@@ -106,6 +111,11 @@ class WorkonHelperTest(cros_test_lib.MockTempDirTestCase):
     # Patch the modules interfaces to the rest of the world.
     self.PatchObject(portage_util, 'FindEbuildForPackage',
                      self._MockFindEbuildForPackage)
+
+    # Assume only versioned-packages is installed.
+    self.PatchObject(
+        portage_util.PortageDB, 'InstalledPackages',
+        return_value=[InstalledPackageMock('sys-apps', 'versioned-package')])
     # This basically turns off behavior related to adding repositories to
     # minilayouts.
     self.PatchObject(git.ManifestCheckout, 'IsFullManifest', return_value=True)
@@ -305,3 +315,9 @@ class WorkonHelperTest(cros_test_lib.MockTempDirTestCase):
     self.assertNotExists(file_path)
     helper.RunCommandInPackages([WORKON_ONLY_ATOM], 'touch %s' % file_name)
     self.assertExists(file_path)
+
+  def testInstalledWorkonAtoms(self):
+    """Test that we can list all the cros workon atoms that are installed."""
+    helper = self.CreateHelper()
+    self.assertEqual(set([VERSIONED_WORKON_ATOM]),
+                     helper.InstalledWorkonAtoms())
diff --git a/scripts/cros_run_unit_tests.py b/scripts/cros_run_unit_tests.py
new file mode 100644
index 000000000..b179121b1
--- /dev/null
+++ b/scripts/cros_run_unit_tests.py
@@ -0,0 +1,103 @@
+# Copyright 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.
+
+"""Tool to run ebuild unittests."""
+
+from __future__ import print_function
+
+import os
+
+from chromite.lib import commandline
+from chromite.lib import chroot_util
+from chromite.lib import cros_build_lib
+from chromite.lib import cros_logging as logging
+from chromite.lib import osutils
+from chromite.lib import workon_helper
+from chromite.lib import portage_util
+
+
+def ParseArgs(argv):
+  """Parse arguments.
+
+  Args:
+    argv: array of arguments passed to the script.
+  """
+  parser = commandline.ArgumentParser(description=__doc__)
+
+  target = parser.add_mutually_exclusive_group(required=True)
+  target.add_argument('--sysroot', type='path', help='Path to the sysroot.')
+  target.add_argument('--board', help='Board name.')
+
+  parser.add_argument('--pretend', default=False, action='store_true',
+                      help='Show the list of packages to be tested and return.')
+  parser.add_argument('--noinstalled_only', dest='installed', default=True,
+                      action='store_false',
+                      help='Test all testable packages, even if they are not '
+                      'currently installed.')
+  parser.add_argument('--package_file', type='path',
+                      help='Path to a file containing the list of packages '
+                      'that should be tested.')
+  parser.add_argument('--blacklist_packages', dest='package_blacklist',
+                      help='Space-separated list of blacklisted packages.')
+  parser.add_argument('--packages',
+                      help='Space-separated list of packages to test.')
+  parser.add_argument('--nowithdebug', action='store_true',
+                      help="Don't build the tests with USE=cros-debug")
+
+  options = parser.parse_args(argv)
+  options.Freeze()
+  return options
+
+
+def main(argv):
+  opts = ParseArgs(argv)
+
+  cros_build_lib.AssertInsideChroot()
+
+  sysroot = opts.sysroot or cros_build_lib.GetSysroot(opts.board)
+  package_blacklist = portage_util.UNITTEST_PACKAGE_BLACKLIST
+  if opts.package_blacklist:
+    package_blacklist |= set(opts.package_blacklist.split())
+
+  packages = set()
+  # The list of packages to test can be passed as a file containing a
+  # space-separated list of package names.
+  # This is used by the builder to test only the packages that were upreved.
+  if opts.package_file and os.path.exists(opts.package_file):
+    packages = set(osutils.ReadFile(opts.package_file).split())
+
+  if opts.packages:
+    packages |= set(opts.packages.split())
+
+  # If no packages were specified, use all testable packages.
+  if not (opts.packages or opts.package_file):
+    workon = workon_helper.WorkonHelper(sysroot)
+    packages = (workon.InstalledWorkonAtoms() if opts.installed
+                else workon.ListAtoms(use_all=True))
+
+  for cp in packages & package_blacklist:
+    logging.info('Skipping blacklisted package %s.', cp)
+
+  packages = packages - package_blacklist
+  pkg_with_test = portage_util.PackagesWithTest(sysroot, packages)
+
+  if packages - pkg_with_test:
+    logging.warning('The following packages do not have tests:')
+    logging.warning('\n'.join(sorted(packages - pkg_with_test)))
+
+  if opts.pretend:
+    print('\n'.join(sorted(pkg_with_test)))
+    return
+
+  env = None
+  if opts.nowithdebug:
+    use_flags = os.environ.get('USE', '')
+    use_flags += ' -cros-debug'
+    env = {'USE': use_flags}
+
+  try:
+    chroot_util.RunUnittests(sysroot, pkg_with_test, extra_env=env)
+  except cros_build_lib.RunCommandError:
+    logging.error('Unittests failed.')
+    raise