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