[autotest] A tool that maps a suite to it's control files.

Currently, this mapping is achieved at runtime, by going
through all control files in search of a suite attribute
that matches the current suite's name. This is an expensive
process, and since the control files don't change after we
freeze them at build time, one that can be avoided with some
preprocessing. This tool will go through all control files
specified in the given root directory and build a suite
to control file map that dynamic_suite can consume at
run time.

BUG=chromium:252398
TEST=Ran the script both stand alone and through emerge
     autotest-all, checked the suite map it generated.

Change-Id: I77df32189039fc9d0f1aa6768a35ae4b361e01ae
Previous-Reviewed-on: https://gerrit.chromium.org/gerrit/61393
(cherry picked from commit 5cb8f91e6853de4ecc351a842eb8f1c1919094bb)
Reviewed-on: https://gerrit.chromium.org/gerrit/63914
Reviewed-by: Alex Miller <milleral@chromium.org>
Commit-Queue: Prashanth Balasubramanian <beeps@chromium.org>
Tested-by: Prashanth Balasubramanian <beeps@chromium.org>
diff --git a/site_utils/control_file_preprocessor.py b/site_utils/control_file_preprocessor.py
new file mode 100755
index 0000000..a039486
--- /dev/null
+++ b/site_utils/control_file_preprocessor.py
@@ -0,0 +1,67 @@
+#!/usr/bin/python
+#
+# Copyright (c) 20123 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Tool for preprocessing control files to build a suite to control files map.
+
+Given an autotest root directory, this tool will bucket tests accroding to
+their suite.Data will be written to stdout (or, optionally a file), eg:
+
+{'suite1': ['path/to/test1/control', 'path/to/test2/control'],
+ 'suite2': ['path/to/test4/control', 'path/to/test5/control']}
+
+This is intended for use only with Chrome OS test suites that leverage the
+dynamic suite infrastructure in server/cros/dynamic_suite.py. It is invoked
+at build time to generate said suite to control files map, which dynamic_suite
+consults at run time to determine which tests belong to a suite.
+"""
+
+
+import collections, json, os, sys
+
+import common
+from autotest_lib.server.cros.dynamic_suite import suite
+from autotest_lib.site_utils import suite_preprocessor
+
+
+def get_suite_control_files(autotest_dir):
+    """
+    Partition all control files in autotest_dir based on suite.
+
+    @param autotest_dir: Directory to walk looking for control files.
+    @return suite_control_files: A dictionary mapping suite->[control files]
+                                 as described in this files docstring.
+    @raise ValueError: If autotest_dir doesn't exist.
+    """
+    if not os.path.exists(autotest_dir):
+      raise ValueError('Could not find directory: %s, failed to map suites to'
+                       ' their control files.' % autotest_dir)
+
+    autotest_dir = autotest_dir.rstrip('/')
+    suite_control_files = collections.defaultdict(list)
+
+    for test in suite_preprocessor.get_all_suite_control_files(autotest_dir):
+        for suite_name in suite.Suite.parse_tag(test.suite):
+            suite_control_files[suite_name].append(
+                test.path.replace('%s/' % autotest_dir, ''))
+    return suite_control_files
+
+
+def main():
+    """
+    Main function.
+    """
+    options = suite_preprocessor.parse_options()
+
+    suite_control_files = get_suite_control_files(options.autotest_dir)
+    if options.output_file:
+        with open(options.output_file, 'w') as file_obj:
+            json.dump(suite_control_files, file_obj)
+    else:
+        print suite_control_files
+
+
+if __name__ == '__main__':
+    sys.exit(main())
diff --git a/site_utils/suite_preprocessor.py b/site_utils/suite_preprocessor.py
index 82ae668..d2a2131 100755
--- a/site_utils/suite_preprocessor.py
+++ b/site_utils/suite_preprocessor.py
@@ -28,6 +28,7 @@
 
 
 def parse_options():
+    """Parse command line arguments."""
     parser = optparse.OptionParser()
     parser.add_option('-a', '--autotest_dir', dest='autotest_dir',
                       default=os.path.abspath(
@@ -42,6 +43,18 @@
     return options
 
 
+def get_all_suite_control_files(autotest_dir):
+    """Find all control files in autotest_dir that have 'SUITE='
+
+    @param autotest_dir: The directory to search for control files.
+    @return: All control files in autotest_dir that have a suite attribute.
+    """
+    fs_getter = Suite.create_fs_getter(autotest_dir)
+    predicate = lambda t: hasattr(t, 'suite')
+    return Suite.find_and_parse_tests(fs_getter, predicate,
+                                      add_experimental=True)
+
+
 def calculate_dependencies(autotest_dir):
     """
     Traverse through the autotest directory and gather together the information
@@ -51,10 +64,8 @@
     @param autotest_dir The path to the autotest directory to examine for tests
     @return A dictionary of the form {suite: {test: [dep, dep]}}.
     """
-    fs_getter = Suite.create_fs_getter(autotest_dir)
-    predicate = lambda t: hasattr(t, 'suite')
     test_deps = {}  #  Format will be {suite: {test: [dep, dep]}}.
-    for test in Suite.find_and_parse_tests(fs_getter, predicate, True):
+    for test in get_all_suite_control_files(autotest_dir):
         for suite in Suite.parse_tag(test.suite):
             suite_deps = test_deps.setdefault(suite, {})
             # Force this to a list so that we can parse it later with
@@ -67,13 +78,14 @@
 
 
 def main():
+    """Main function."""
     options = parse_options()
 
     test_deps = calculate_dependencies(options.autotest_dir)
 
     if options.output_file:
-        with open(options.output_file, 'w+') as fd:
-            fd.write('%r' % test_deps)
+        with open(options.output_file, 'w') as file_obj:
+            file_obj.write('%r' % test_deps)
     else:
         print '%r' % test_deps