1#!/usr/bin/python2 2# 3# Copyright (c) 20123 The Chromium OS Authors. All rights reserved. 4# Use of this source code is governed by a BSD-style license that can be 5# found in the LICENSE file. 6 7"""Tool for preprocessing control files to build a suite to control files map. 8 9Given an autotest root directory, this tool will bucket tests accroding to 10their suite.Data will be written to stdout (or, optionally a file), eg: 11 12{'suite1': ['path/to/test1/control', 'path/to/test2/control'], 13 'suite2': ['path/to/test4/control', 'path/to/test5/control']} 14 15This is intended for use only with Chrome OS test suites that leverage the 16dynamic suite infrastructure in server/cros/dynamic_suite.py. It is invoked 17at build time to generate said suite to control files map, which dynamic_suite 18consults at run time to determine which tests belong to a suite. 19""" 20 21 22from __future__ import absolute_import 23from __future__ import division 24from __future__ import print_function 25 26import collections, json, os, sys 27 28import common 29from autotest_lib.server.cros.dynamic_suite import suite 30from autotest_lib.site_utils import suite_preprocessor 31 32 33# A set of SUITES that we choose not to preprocess as they might have tests 34# added later. 35SUITE_DENYLIST = set(['au']) 36 37 38def _get_control_files_to_process(autotest_dir): 39 """Find all control files in autotest_dir that have 'SUITE=' 40 41 @param autotest_dir: The directory to search for control files. 42 @return: All control files in autotest_dir that have a suite attribute. 43 """ 44 fs_getter = suite.Suite.create_fs_getter(autotest_dir) 45 predicate = lambda t: hasattr(t, 'suite') 46 return suite.Suite.find_and_parse_tests(fs_getter, predicate, 47 add_experimental=True) 48 49 50def get_suite_control_files(autotest_dir, external_autotest_dirs=None): 51 """ 52 Partition all control files in autotest_dir based on suite. 53 54 @param autotest_dir: Directory to walk looking for control files. 55 @param external_autotest_dirs: A list of directories under which to search 56 for extra Autotest tests. Defaults to None. 57 58 @return suite_control_files: A dictionary mapping suite->[control files] 59 as described in this files docstring. 60 @raise ValueError: If autotest_dir doesn't exist. 61 """ 62 if not os.path.exists(autotest_dir): 63 raise ValueError('Could not find directory: %s, failed to map suites to' 64 ' their control files.' % autotest_dir) 65 66 suite_control_files = collections.defaultdict(list) 67 for d in [autotest_dir] + (external_autotest_dirs or []): 68 d = d.rstrip('/') 69 for test in _get_control_files_to_process(d): 70 for suite_name in test.suite_tag_parts: 71 if suite_name in SUITE_DENYLIST: 72 continue 73 74 suite_control_files[suite_name].append( 75 test.path.replace('%s/' % d, '')) 76 return suite_control_files 77 78 79def main(): 80 """ 81 Main function. 82 """ 83 options = suite_preprocessor.parse_options() 84 85 if options.extra_autotest_dirs: 86 extra_autotest_dirs = options.extra_autotest_dirs.split(',') 87 else: 88 extra_autotest_dirs = None 89 90 suite_control_files = get_suite_control_files(options.autotest_dir, 91 extra_autotest_dirs) 92 if options.output_file: 93 with open(options.output_file, 'w') as file_obj: 94 json.dump(suite_control_files, file_obj) 95 else: 96 print(suite_control_files) 97 98 99if __name__ == '__main__': 100 sys.exit(main()) 101