Home | History | Annotate | Download | only in site_utils 
      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:
     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 
     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
     95     elif not hasattr(cd, 'suite') and cd_attrs:
     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())
    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 
    136 if __name__ == '__main__':
    137   sys.exit(main(sys.argv[1:]))
    138