Add module discovery and autoloading to expect_tests.

scripts/slave/unittests/expect_tests/main.py

Index: scripts/slave/unittests/expect_tests/main.py
diff --git a/scripts/slave/unittests/expect_tests/main.py b/scripts/slave/unittests/expect_tests/main.py
index 5a266b1136eea7e6caff7030d0a86a12a56c3bfe..ce9f11964bd1a3e7aa6dc0788ad6b4f925b29605 100644
--- a/scripts/slave/unittests/expect_tests/main.py
+++ b/scripts/slave/unittests/expect_tests/main.py
@@ -4,6 +4,7 @@
 
 import argparse
 import multiprocessing
+import pkgutil
 import sys
 
 from .cover import CoverageContext
@@ -12,6 +13,11 @@ 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,
@@ -35,8 +41,8 @@ class _test_completer(object):
       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 = []
@@ -45,12 +51,13 @@ class _test_completer(object):
         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
@@ -83,6 +90,11 @@ def _parse_args(args, test_gen):
           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)'
@@ -99,7 +111,7 @@ def _parse_args(args, test_gen):
             '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)
 
@@ -121,7 +133,32 @@ def _parse_args(args, test_gen):
   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:
@@ -138,23 +175,28 @@ def main(name, test_gen, cover_branches=False, args=None):
     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)