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 6d5ddbf5b98211b931d0d19dfcc434cea8ea7092..3e4e27e9f08f95f866035738403295b689173716 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 inspect
 import logging
@@ -15,7 +16,7 @@ from cStringIO import StringIO
 
 from .type_definitions import (
     Test, UnknownError, TestError, NoMatchingTestsError, MultiTest,
-    Result, ResultStageAbort)
+    Result, ResultStageAbort, Cleanup)
 
 
 class ResetableStringIO(object):
@@ -29,15 +30,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
@@ -57,53 +58,69 @@ 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:
+        if hasattr(gen, '_covers'):
+          # decorated with expect_tests.covers
+          gen_cover = gen._covers()  # pylint: disable=W0212
+        else:
+          gen_cover = [inspect.getabsfile(gen)]
+        gen_cover_ctx = cover_ctx(include=gen_cover)
 
-      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:
@@ -177,7 +194,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()
@@ -190,17 +207,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:
-    if hasattr(test_gen, '_covers'):
-      # decorated with expect_tests.covers
-      gen_cover = test_gen._covers()  # pylint: disable=W0212
-    else:
-      gen_cover = [inspect.getabsfile(test_gen)]
-    gen_cover_ctx = cover_ctx(include=gen_cover)
-
   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:
@@ -226,7 +234,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)

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

em... so func_call is crossing process boundaries?

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

yep. I'm using this to delete the temporary directory (best effort) that the
test generator created.

I'm wondering if 'temporary directory' should be a service that expect_tests
just provides...

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

I hope it's clearly marked in docstring that cleanup function is executing in
another process? Since fork happens early, this clean up function may see stale
values of globals, etc. It may be confusing.

If it's only for temp dirs, then I personally prefer "service" approach.

+            else:
+              yield obj
           except Queue.Empty:
             break