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 .pipeline import result_loop
 
+from .unittest_helper import _is_unittest, UnittestTestCase
+
+
+EVERYTHING = object()

[Vadim Sh., 2014/06/27 21:40:08]

It's a not very descriptive name. every what?

ALL_MODULES?

[iannucci, 2014/06/28 16:22:17]

Done.

+
 
 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.'
+      )

[Vadim Sh., 2014/06/27 21:40:08]

nit: keep ) on previous line
or move ) on line 91 on a next line

I hate formatting inconsistencies within 10 line of code >_< 

[iannucci, 2014/06/28 16:22:17]

Done.

+
     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 EVERYTHING:
+    # if we're running directly
+    if __name__ == '__main__' or test_modules is EVERYTHING:
+      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 getattr(obj, '_expect_tests_generator', False):
+        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
 
     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():

[Vadim Sh., 2014/06/27 21:40:08]

You seemed to converge on convention to use CapitalCase for generators and
lower_case_and_underscores for regular functions?.. I'm personally don't like
this :) Just saying.

[iannucci, 2014/06/28 16:22:17]

I think it's just inconsistent. The idea was CamelCase for (exported from
module) and snake_case for everything else.

I can rename. Keep in mind, this is inside a docstring, even though it doesn't
look like it :p

[pgervais, 2014/06/30 13:56:03]

Don't we have a style guide for this? I really don't like the added mental load
of knowing if the function is exported or before knowing how to write it. I
consider this to be a mistake in Go, please don't get this into our Python
codebase as well.

[Vadim Sh., 2014/06/30 17:15:46]

+1 to keeping single naming convention for public and private functions

       yield expect_tests.Test('happy', happy_fn, args=('happy',))
 
     if __name__ == '__main__':
       expect_tests.main('happy_test_suite', Gen)
 
-  @param name: Name of the test suite.
-  @param test_gen: A Generator which yields Test objects.
+  @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 cover_branches: Include branch coverage data (rather than just line
                          coverage)
   @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