Add module discovery and autoloading to expect_tests.

scripts/slave/unittests/expect_tests/main.py

 # Copyright 2014 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 import argparse
 import multiprocessing
+import pkgutil
 import sys
 
 from .cover import CoverageContext
 
-from . import handle_list, handle_debug, handle_train, handle_test
+from . import handle_list, handle_debug, handle_train, handle_test, util
 
 from .pipeline import result_loop
 
+from .unittest_helper import _is_unittest, UnittestTestCase
+
+
+ALL_MODULES = object()
+
 
 HANDLERS = {
   'list': handle_list.ListHandler,
   'debug': handle_debug.DebugHandler,
   'train': handle_train.TrainHandler,
   'test': handle_test.TestHandler,
 }
 
 
 class _test_completer(object):
   """Implements the argcomplete completer interface for the test_glob command
   line argument.
 
   See: https://pypi.python.org/pypi/argcomplete
 
   This is automatically wired up if you have enabled bash completion in the
   infra repo: https://chromium.googlesource.com/infra/infra
   """
   class FakeOptions(object):
     def __init__(self, **kwargs):
       for k, v in kwargs.iteritems():
         setattr(self, k, v)
 
-  def __init__(self, gen):
-    self._gen = gen
+  def __init__(self, test_modules):
+    self._test_modules = test_modules
 
   def __call__(self, prefix, **_):
     handle_list.ListHandler.COMPLETION_LIST = []
     options = self.FakeOptions(
         handler=handle_list.ListHandler,
         test_glob=[prefix],
         jobs=1,
     )
-    ctx = CoverageContext('', [], [], False, None, None, False)
-    result_loop(self._gen, ctx.create_subprocess_context(), options)
+    ctx = CoverageContext(False, False, False)
+    test_gens = get_test_gens(self._test_modules)
+    result_loop(test_gens, ctx.create_subprocess_context(), options)
     return handle_list.ListHandler.COMPLETION_LIST
 
 
-def _parse_args(args, test_gen):
+def _parse_args(args, test_modules):
   args = args or sys.argv[1:]
 
   # Set the default mode if not specified and not passing --help
   search_names = set(HANDLERS.keys() + ['-h', '--help'])
   if not any(arg in search_names for arg in args):
     args.insert(0, 'test')
 
   parser = argparse.ArgumentParser()
   subparsers = parser.add_subparsers(
       title='Mode (default "test")', dest='mode',
@@ skipping 12 matching lines @@
         help='be quiet (only print failures)')
     mg.add_argument(
         '--verbose', action='store_true', help='be verbose')
 
     if not h.SKIP_RUNLOOP:
       sp.add_argument(
           '--jobs', metavar='N', type=int,
           default=multiprocessing.cpu_count(),
           help='run N jobs in parallel (default %(default)s)')
 
+      sp.add_argument(
+          '--force_coverage', action='store_true',
+          help='Enable coverage report even when specifying a test filter.')
+
     sp.add_argument(
         '--test_list', metavar='FILE',
         help='take the list of test globs from the FILE (use "-" for stdin)'
     ).completer = lambda **_: []
 
     sp.add_argument(
         '--html_report', metavar='DIR',
         help='directory to write html report (default: disabled)'
     ).completer = lambda **_: []
 
     sp.add_argument(
         'test_glob', nargs='*', help=(
             'glob to filter the tests acted on. If the glob begins with "-" '
             'then it acts as a negation glob and anything which matches it '
             'will be skipped. If a glob doesn\'t have "*" in it, "*" will be '
             'implicitly appended to the end')
-    ).completer = _test_completer(test_gen)
+    ).completer = _test_completer(test_modules)
 
   opts = parser.parse_args(args)
 
   if not hasattr(opts, 'jobs'):
     opts.jobs = 0
   elif opts.jobs < 1:
     parser.error('--jobs was less than 1')
 
   if opts.test_list:
     fh = sys.stdin if opts.test_list == '-' else open(opts.test_list, 'rb')
     with fh as tl:
       opts.test_glob += [l.strip() for l in tl.readlines()]
 
   opts.handler = HANDLERS[opts.mode]
 
   del opts.test_list
   del opts.mode
 
   return opts
 
 
-def main(name, test_gen, cover_branches=False, args=None):
+def get_test_gens(test_modules):
+  test_gens = []
+  if not test_modules or test_modules is ALL_MODULES:
+    # if we're running directly
+    if __name__ == '__main__' or test_modules is ALL_MODULES:
+      test_modules = []
+      for importer, modname, ispkg in pkgutil.walk_packages(path=['.']):
+        if not ispkg and modname.endswith('_test'):
+          if modname in sys.modules:
+            test_modules.append(sys.modules[modname])
+          else:
+            test_modules.append(
+                importer.find_module(modname).load_module(modname))
+    else:  # a wrapper main() script
+      test_modules = [sys.modules['__main__']]
+  for mod in test_modules:
+    for obj in mod.__dict__.values():
+      if util.is_test_generator(obj):
+        test_gens.append(obj)
+      elif _is_unittest(obj):
+        test_gens.append(UnittestTestCase(obj))
+  return test_gens
+
+
+# TODO(iannucci): have Test determine cover_branches
+def main(cover_branches=False, test_modules=None, args=None):
   """Entry point for tests using expect_tests.
 
   Example:
-    import expect_tests
+  >>> import expect_tests
 
-    def happy_fn(val):
-      # Usually you would return data which is the result of some deterministic
-      # computation.
-      return expect_tests.Result({'neet': '%s string value' % val})
+  >>> def happy_fn(val):
+  >>>   # Usually you would return data which is the result of some
+  >>>   #deterministic computation.
+  >>>   return expect_tests.Result({'neet': '%s string value' % val})
 
-    def Gen():
-      yield expect_tests.Test('happy', happy_fn, args=('happy',))
+  >>> @expect_tests.test_generator
+  >>> def gen():
+  >>>   yield expect_tests.Test(
+  >>>     __package__ + '.happy',
+  >>>     expect_tests.FuncCall(happy_fn, 'happy'))
 
-    if __name__ == '__main__':
-      expect_tests.main('happy_test_suite', Gen)
+  >>> if __name__ == '__main__':
+  >>>   expect_tests.main(cover_branches=True)
 
-  @param name: Name of the test suite.
-  @param test_gen: A Generator which yields Test objects.
   @param cover_branches: Include branch coverage data (rather than just line
                          coverage)
+  @param test_modules: Modules containing expect_tests generators and/or
+                       unittests. Defaults to the __main__ module, or if
+                       this script is invoked directly, all '_test' modules
+                       under the current working directory.
   @param args: Commandline args (starting at argv[1])
   """
   try:
-    opts = _parse_args(args, test_gen)
+    opts = _parse_args(args, test_modules)
 
-    cover_ctx = CoverageContext(name, cover_branches, opts.html_report,
+    cover_ctx = CoverageContext(cover_branches, opts.html_report,
                                 not opts.handler.SKIP_RUNLOOP)
 
+    with cover_ctx.create_subprocess_context():
+      test_gens = get_test_gens(test_modules)
+
     error, killed = result_loop(
-        test_gen, cover_ctx.create_subprocess_context(), opts)
+        test_gens, cover_ctx.create_subprocess_context(), opts)
 
     cover_ctx.cleanup()
-    if not killed and not opts.test_glob:
+    if not killed and (opts.force_coverage or not opts.test_glob):
       if not cover_ctx.report(opts.verbose):
         sys.exit(2)
 
     sys.exit(error or killed)
   except KeyboardInterrupt:
     pass