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
8ATTRIBUTES either not match to SUITE or not in the attribute whitelist."""
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', 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:
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)
45    count = len(changelist)
46
47    logging.info('Starting to seed attributes in %d control files...' % count)
48    # Modify attributes based on suite for the control files not match.
49    for path in changelist:
50      logging.info('Seeding ATTRIBUTES in %s' % path)
51      count = count - 1
52      logging.info('%d files remaining...' % count)
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
63def 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
95    elif not hasattr(cd, 'suite') and cd_attrs:
96      unmatch_pathlist.append(path)
97
98  return unmatch_pathlist
99
100
101def 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())
118  attr_items = list(attr_items)
119  attr_items.sort(key = str.lower)
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()
126    index = [i for i, val in enumerate(lines) if val.startswith('SUITE =') or
127             val.startswith('SUITE=')][0]
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
136if __name__ == '__main__':
137  sys.exit(main(sys.argv[1:]))
138