Add module discovery and autoloading to expect_tests.

scripts/slave/unittests/expect_tests/pipeline.py

Index: scripts/slave/unittests/expect_tests/pipeline.py
diff --git a/scripts/slave/unittests/expect_tests/pipeline.py b/scripts/slave/unittests/expect_tests/pipeline.py
index 9f23d3427bb64bd273304277308c276a851acb06..8ebb4f62d514bb5f63f6063076858fd922276104 100644
--- a/scripts/slave/unittests/expect_tests/pipeline.py
+++ b/scripts/slave/unittests/expect_tests/pipeline.py
@@ -3,6 +3,7 @@
 # found in the LICENSE file.
 
 import Queue
+import atexit
 import glob
 import logging
 import multiprocessing
@@ -14,7 +15,7 @@ from cStringIO import StringIO
 
 from .type_definitions import (
     Test, UnknownError, TestError, NoMatchingTestsError, MultiTest,
-    Result, ResultStageAbort)
+    Result, ResultStageAbort, Cleanup)
 
 from . import util
 
@@ -30,15 +31,15 @@ class ResetableStringIO(object):
     return getattr(self._stream, key)
 
 
-def gen_loop_process(gen, test_queue, result_queue, opts, kill_switch,
+def gen_loop_process(gens, test_queue, result_queue, opts, kill_switch,
                      cover_ctx):
-  """Generate `Test`'s from |gen|, and feed them into |test_queue|.
+  """Generate `Test`'s from |gens|, and feed them into |test_queue|.
 
   Non-Test instances will be translated into `UnknownError` objects.
 
   On completion, feed |opts.jobs| None objects into |test_queue|.
 
-  @param gen: generator yielding Test() instances.
+  @param gens: list of generators yielding Test() instances.
   @type test_queue: multiprocessing.Queue()
   @type result_queue: multiprocessing.Queue()
   @type opts: argparse.Namespace
@@ -58,53 +59,64 @@ def gen_loop_process(gen, test_queue, result_queue, opts, kill_switch,
       '^%s$' % '|'.join('(?:%s)' % glob.fnmatch.translate(g[1:])
                         for g in globs if g[0] == '-'))
 
+  SENTINEL = object()
+
   def generate_tests():
     paths_seen = set()
     seen_tests = False
     try:
-      with cover_ctx:
-        gen_inst = gen()
+      for gen in gens:
+        gen_cover_ctx = cover_ctx(include=util.get_cover_list(gen))
 
-      while not kill_switch.is_set():
-        with cover_ctx:
-          root_test = next(gen_inst)
+        with gen_cover_ctx:
+          gen_inst = gen()
 
-        if kill_switch.is_set():
-          break
+        while not kill_switch.is_set():
+          with gen_cover_ctx:
+            root_test = next(gen_inst, SENTINEL)
+
+          if root_test is SENTINEL:
+            break
 
-        ok_tests = []
+          if kill_switch.is_set():
+            break
 
-        if isinstance(root_test, MultiTest):
-          subtests = root_test.tests
-        else:
-          subtests = [root_test]
+          ok_tests = []
 
-        for subtest in subtests:
-          if not isinstance(subtest, Test):
-            result_queue.put_nowait(
-                UnknownError('Got non-[Multi]Test isinstance from generator: %r'
-                             % subtest))
+          if isinstance(root_test, Cleanup):
+            result_queue.put_nowait(root_test)
             continue
 
-          test_path = subtest.expect_path()
-          if test_path is not None and test_path in paths_seen:
-            result_queue.put_nowait(
-                TestError(subtest, 'Duplicate expectation path!'))
+          if isinstance(root_test, MultiTest):
+            subtests = root_test.tests
           else:
-            if test_path is not None:
-              paths_seen.add(test_path)
-            name = subtest.name
-            if not neg_matcher.match(name) and matcher.match(name):
-              ok_tests.append(subtest)
-
-        if ok_tests:
-          seen_tests = True
-          yield root_test.restrict(ok_tests)
+            subtests = [root_test]
+
+          for subtest in subtests:
+            if not isinstance(subtest, Test):
+              result_queue.put_nowait(
+                  UnknownError(
+                      'Got non-[Multi]Test isinstance from generator: %r'
+                      % subtest))
+              continue
+
+            test_path = subtest.expect_path()
+            if test_path is not None and test_path in paths_seen:
+              result_queue.put_nowait(
+                  TestError(subtest, 'Duplicate expectation path!'))
+            else:
+              if test_path is not None:
+                paths_seen.add(test_path)
+              name = subtest.name
+              if not neg_matcher.match(name) and matcher.match(name):
+                ok_tests.append(subtest)
+
+          if ok_tests:
+            seen_tests = True
+            yield root_test.restrict(ok_tests)
 
       if not seen_tests:
         result_queue.put_nowait(NoMatchingTestsError())
-    except StopIteration:
-      pass
     except KeyboardInterrupt:
       pass
     finally:
@@ -178,7 +190,7 @@ def run_loop_process(test_queue, result_queue, opts, kill_switch, cover_ctx):
                               result_queue.put_nowait)
 
 
-def result_loop(test_gen, cover_ctx, opts):
+def result_loop(test_gens, cover_ctx, opts):
   kill_switch = multiprocessing.Event()
   def handle_killswitch(*_):
     kill_switch.set()
@@ -191,12 +203,8 @@ def result_loop(test_gen, cover_ctx, opts):
   test_queue = multiprocessing.Queue()
   result_queue = multiprocessing.Queue()
 
-  gen_cover_ctx = cover_ctx
-  if cover_ctx.enabled:
-    gen_cover_ctx = cover_ctx(include=util.get_cover_list(test_gen))
-
   test_gen_args = (
-      test_gen, test_queue, result_queue, opts, kill_switch, gen_cover_ctx)
+      test_gens, test_queue, result_queue, opts, kill_switch, cover_ctx)
 
   procs = []
   if opts.handler.SKIP_RUNLOOP:
@@ -222,7 +230,11 @@ def result_loop(test_gen, cover_ctx, opts):
       while not kill_switch.is_set():
         while not kill_switch.is_set():
           try:
-            yield result_queue.get(timeout=0.1)
+            obj = result_queue.get(timeout=0.1)
+            if isinstance(obj, Cleanup):
+              atexit.register(obj.func_call)
+            else:
+              yield obj
           except Queue.Empty:
             break