Shuqian Zhao | 12f0915 | 2015-04-14 12:56:35 -0700 | [diff] [blame] | 1 | #!/usr/bin/python |
| 2 | |
| 3 | # Copyright (c) 2015 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 | """Method to add or modify ATTRIBUTES in the test control files whose |
| 8 | ATTRIBUTES either not match to SUITE or not in the attribute whitelist.""" |
| 9 | |
| 10 | import argparse |
| 11 | import logging |
| 12 | import os |
| 13 | import sys |
| 14 | |
| 15 | import common |
| 16 | from autotest_lib.client.common_lib import control_data |
| 17 | from autotest_lib.server.cros.dynamic_suite.suite import Suite |
| 18 | |
| 19 | |
| 20 | def main(argv): |
| 21 | """main scripts to seed attributes in test control files. |
| 22 | |
| 23 | Args: |
| 24 | @param argv: Command line arguments including `sys.argv[0]`. |
| 25 | """ |
| 26 | # Parse execution cmd |
| 27 | parser = argparse.ArgumentParser( |
| 28 | description='Seed ATTRIBUTES in test control files.') |
| 29 | parser.add_argument('--execute', action='store_true', default=False, |
| 30 | help='Execute the script to seed attributes in all ' |
| 31 | 'test control files.') |
| 32 | args = parser.parse_args(argv) |
| 33 | |
| 34 | # When execute is True, run the script to seed attributes in control files. |
| 35 | if args.execute: |
Shuqian Zhao | 12f0915 | 2015-04-14 12:56:35 -0700 | [diff] [blame] | 36 | # Get the whitelist path, hardcode the path currently |
| 37 | path_whitelist = os.path.join(common.autotest_dir, |
| 38 | 'site_utils/attribute_whitelist.txt') |
| 39 | |
| 40 | # Go through all control file, check whether attribute matches suite. Return |
| 41 | # a changelist which contains the paths to the control files not match. |
| 42 | fs_getter = Suite.create_fs_getter(common.autotest_dir) |
| 43 | changelist = AttrSuiteMatch( |
| 44 | fs_getter.get_control_file_list(), path_whitelist) |
Shuqian Zhao | d71d2a9 | 2015-04-29 14:46:10 -0700 | [diff] [blame] | 45 | count = len(changelist) |
Shuqian Zhao | 12f0915 | 2015-04-14 12:56:35 -0700 | [diff] [blame] | 46 | |
Shuqian Zhao | d71d2a9 | 2015-04-29 14:46:10 -0700 | [diff] [blame] | 47 | logging.info('Starting to seed attributes in %d control files...' % count) |
Shuqian Zhao | 12f0915 | 2015-04-14 12:56:35 -0700 | [diff] [blame] | 48 | # Modify attributes based on suite for the control files not match. |
| 49 | for path in changelist: |
Shuqian Zhao | d71d2a9 | 2015-04-29 14:46:10 -0700 | [diff] [blame] | 50 | logging.info('Seeding ATTRIBUTES in %s' % path) |
| 51 | count = count - 1 |
| 52 | logging.info('%d files remaining...' % count) |
Shuqian Zhao | 12f0915 | 2015-04-14 12:56:35 -0700 | [diff] [blame] | 53 | SeedAttributes(path) |
| 54 | |
| 55 | logging.info('Finished seeding attributes.') |
| 56 | |
| 57 | # When not specify 'execute' in cmd, not modify control files. |
| 58 | else: |
| 59 | logging.info('No files are modified. To seed attributes in control files, ' |
| 60 | 'please add \'--execute\' argument when run the script.') |
| 61 | |
| 62 | |
| 63 | def AttrSuiteMatch(path_list, path_whitelist): |
| 64 | """Check whether attributes are in the attribute whitelist and match with the |
| 65 | suites in the control files. |
| 66 | |
| 67 | Args: |
| 68 | @param path_list: a list of path to the control files to be checked. |
| 69 | @param path_whitelist: path to the attribute whitelist. |
| 70 | |
| 71 | Returns: |
| 72 | A list of paths to the control files that failed at checking. |
| 73 | """ |
| 74 | unmatch_pathlist = [] |
| 75 | |
| 76 | # Read the whitelist to a set, if path is invalid, throw IOError. |
| 77 | with open(path_whitelist, 'r') as f: |
| 78 | whitelist = {line.strip() for line in f.readlines() if line.strip()} |
| 79 | |
| 80 | # Read the attr in the control files, check with whitelist and suite. |
| 81 | for path in path_list: |
| 82 | cd = control_data.parse_control(path, True) |
| 83 | cd_attrs = cd.attributes |
| 84 | |
| 85 | # Test whether attributes in the whitelist |
| 86 | if not (whitelist >= cd_attrs): |
| 87 | unmatch_pathlist.append(path) |
| 88 | # Test when suite exists, whether attributes match suites |
| 89 | if hasattr(cd, 'suite'): |
| 90 | target_attrs = set( |
| 91 | 'suite:' + x.strip() for x in cd.suite.split(',') if x.strip()) |
| 92 | if cd_attrs != target_attrs: |
| 93 | unmatch_pathlist.append(path) |
| 94 | # Test when suite not exists, whether attributes is empty |
Shuqian Zhao | d71d2a9 | 2015-04-29 14:46:10 -0700 | [diff] [blame] | 95 | elif not hasattr(cd, 'suite') and cd_attrs: |
Shuqian Zhao | 12f0915 | 2015-04-14 12:56:35 -0700 | [diff] [blame] | 96 | unmatch_pathlist.append(path) |
| 97 | |
| 98 | return unmatch_pathlist |
| 99 | |
| 100 | |
| 101 | def SeedAttributes(path_controlfile): |
| 102 | """Seed attributes in a control file. |
| 103 | |
| 104 | Read and re-write a control file with modified contents with attributes added. |
| 105 | |
| 106 | Args: |
| 107 | @param path_controlfile: path to control file |
| 108 | |
| 109 | Returns: |
| 110 | None |
| 111 | """ |
| 112 | # Parse attribute from suite, and prepare ATTRIBUTES line. |
| 113 | cd = control_data.parse_control(path_controlfile, True) |
| 114 | suite = cd.suite |
| 115 | |
| 116 | attr_items = set( |
| 117 | 'suite:' + x.strip() for x in suite.split(',') if x.strip()) |
Shuqian Zhao | aa7ec8c | 2015-05-06 11:11:58 -0700 | [diff] [blame] | 118 | attr_items = list(attr_items) |
| 119 | attr_items.sort(key = str.lower) |
Shuqian Zhao | 12f0915 | 2015-04-14 12:56:35 -0700 | [diff] [blame] | 120 | attr_line = ', '.join(attr_items) |
| 121 | attr_line = 'ATTRIBUTES = \"' + attr_line + '\"\n' |
| 122 | |
| 123 | # Read control file and modify the suite line with attribute added. |
| 124 | with open(path_controlfile, 'r') as f: |
| 125 | lines = f.readlines() |
Shuqian Zhao | d71d2a9 | 2015-04-29 14:46:10 -0700 | [diff] [blame] | 126 | index = [i for i, val in enumerate(lines) if val.startswith('SUITE =') or |
| 127 | val.startswith('SUITE=')][0] |
Shuqian Zhao | 12f0915 | 2015-04-14 12:56:35 -0700 | [diff] [blame] | 128 | suite_line = lines[index] |
| 129 | lines[index] = attr_line + suite_line |
| 130 | |
| 131 | # Write the modified contents back to file |
| 132 | with open(path_controlfile, 'w') as f: |
| 133 | f.writelines(lines) |
| 134 | |
| 135 | |
| 136 | if __name__ == '__main__': |
| 137 | sys.exit(main(sys.argv[1:])) |