[crostestutils] Add CLI to tauto test metadata modifier script

Add a CLI to specify tests on which to filter and actions to perform.

Also add appending/prepending contacts functionality and fix a bug
where only the first successful action would be applied to a file.

BUG=None
TEST=unit tests pass, as does local verification

Change-Id: If062afa9d627cfc5595549e95fc40dc37b0d023e
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/crostestutils/+/4752955
Commit-Queue: Katherine Threlkeld <kathrelkeld@chromium.org>
Tested-by: Katherine Threlkeld <kathrelkeld@chromium.org>
Reviewed-by: Jesse McGuire <jessemcguire@google.com>
diff --git a/metadata_modifier/tauto_modify/actions.py b/metadata_modifier/tauto_modify/actions.py
index b3549c4..d24c6f7 100644
--- a/metadata_modifier/tauto_modify/actions.py
+++ b/metadata_modifier/tauto_modify/actions.py
@@ -11,7 +11,7 @@
     """Return an action which removes the given list of emails from 'contacts'.
 
     Args:
-        emails: a list of strings, e.g. ['foo@google.com']
+        emails: A list of string emails, e.g. ['foo@google.com'].
 
     Returns:
         An action function that acts on a ControlFile and returns a boolean.
@@ -26,3 +26,43 @@
                 modified = True
         return modified
     return output
+
+
+def append_contacts(emails):
+    """Return an action which appends the given emails to 'contacts'.
+
+    Args:
+        emails: A list of string emails, e.g. ['foo@google.com'].
+
+    Returns:
+        An action function that acts on a ControlFile and returns a boolean.
+    """
+    def output(cf):
+        if 'contacts' not in cf.metadata:
+            return False
+        for email in emails:
+            if email in cf.metadata['contacts']:
+                cf.metadata['contacts'].remove(email)
+        cf.metadata['contacts'] += emails
+        return True
+    return output
+
+
+def prepend_contacts(emails):
+    """Return an action which prepends the given emails to 'contacts'.
+
+    Args:
+        emails: A list of string emails, e.g. ['foo@google.com'].
+
+    Returns:
+        An action function that acts on a ControlFile and returns a boolean.
+    """
+    def output(cf):
+        if 'contacts' not in cf.metadata:
+            return False
+        for email in emails:
+            if email in cf.metadata['contacts']:
+                cf.metadata['contacts'].remove(email)
+        cf.metadata['contacts'] = emails + cf.metadata['contacts']
+        return True
+    return output
diff --git a/metadata_modifier/tauto_modify/cf_parse.py b/metadata_modifier/tauto_modify/cf_parse.py
index da0b3f3..c0ddbf2 100644
--- a/metadata_modifier/tauto_modify/cf_parse.py
+++ b/metadata_modifier/tauto_modify/cf_parse.py
@@ -141,3 +141,4 @@
         self.contents = (self.contents[:self.metadata_start] +
                          new_metadata +
                          self.contents[self.metadata_end:])
+        self.metadata_end = self.metadata_start + len(new_metadata)
diff --git a/metadata_modifier/tauto_modify/tauto_modify.py b/metadata_modifier/tauto_modify/tauto_modify.py
index d596068..215e1fa 100644
--- a/metadata_modifier/tauto_modify/tauto_modify.py
+++ b/metadata_modifier/tauto_modify/tauto_modify.py
@@ -4,8 +4,9 @@
 
 """Helper script to parse autotest control files for data gathering."""
 
-import os
+import argparse
 import pathlib
+import re
 
 import actions
 import cf_parse
@@ -13,29 +14,50 @@
 
 
 # ChromeOS src/ dir relative to this file.
-SRC_DIR = pathlib.Path(__file__).joinpath('../../../../..').resolve()
+DEFAULT_SRC_DIR = pathlib.Path(__file__).joinpath('../../../../..').resolve()
 
-def modify_control_files(action_func_list, filter_func_list, dry_run):
-    # Places to look for control files, relative to SRC_DIR.
-    AUTOTEST_DIRS = [
-            'third_party/autotest/files/client/site_tests/',
-            'third_party/autotest/files/server/site_tests/',
-            'third_party/autotest-private/client/site_tests/',
-    ]
+def modify_control_files(src_dir, actions_list, filters_list, dry_run=True,
+                         public=True, private=True, client=True, server=True):
+    """Apply the given actions to control files if they pass the given filters.
 
-    for tests_dir in [pathlib.Path(SRC_DIR, d) for d in AUTOTEST_DIRS]:
+    Args:
+        src_dir: Relative path to the chromium.org src/ directory.
+        actions_list: List of action functions applied to a ControlFile object.
+        filters_list: List of filter functions applied to a ControlFile object.
+        dry_run: True if no files should be actually modified, just printed.
+        public: True if public tests should be searched.
+        private: True if private tests should be searched.
+        client: True if client tests should be searched.
+        server: True if server tests should be searched.
+    """
+    # Places to look for control files, relative to src_dir.
+    PUBLIC_CLIENT_DIR = 'third_party/autotest/files/client/site_tests/'
+    PUBLIC_SERVER_DIR = 'third_party/autotest/files/server/site_tests/'
+    PRIVATE_CLIENT_DIR = 'third_party/autotest-private/client/site_tests/'
+
+    autotest_dirs = []
+    if public:
+        if client:
+            autotest_dirs.append(PUBLIC_CLIENT_DIR)
+        if server:
+            autotest_dirs.append(PUBLIC_SERVER_DIR)
+    if private and client:
+        autotest_dirs.append(PRIVATE_CLIENT_DIR)
+
+    for tests_dir in [pathlib.Path(src_dir, d) for d in autotest_dirs]:
         for cf_path in tests_dir.glob('*/control*'):
             cf = cf_parse.ControlFile(cf_path)
             if not cf.is_valid:
                 continue
 
             # Skip this control file if it doesn't match all the given filters.
-            if not all(filter_func(cf) for filter_func in filter_func_list):
+            if not all(filter_func(cf) for filter_func in filters_list):
                 continue
 
             # Apply the given actions to this control file.
-            if not any(action_func(cf) for action_func in action_func_list):
-                continue
+            modified = False
+            for action_func in actions_list:
+                modified = action_func(cf) or modified
             cf.update_contents()
 
             if dry_run:
@@ -47,12 +69,96 @@
                     f.write(cf.contents)
 
 
+def _set_up_args():
+    """Define CLI arguments."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--src_dir', default=DEFAULT_SRC_DIR,
+                        help=('Path to the top-level chromiumos src/ directory '
+                              'in your repo. Defaults to the path relative to '
+                              'this file\'s location.'))
+    parser.add_argument('--write_out', action='store_true',
+                        help=('Write out changes to your local repo. Please '
+                              'do a dry run first!'))
+    public_private_group = parser.add_mutually_exclusive_group()
+    public_private_group.add_argument(
+            '--public_only', action='store_true',
+            help='Limit filtering to only public autotest tests.')
+    public_private_group.add_argument(
+            '--private_only', action='store_true',
+            help='Limit filtering to only autotest-private tests.')
+    client_server_group = parser.add_mutually_exclusive_group()
+    client_server_group.add_argument(
+            '--client_only', action='store_true',
+            help='Limit filtering to only client tests.')
+    client_server_group.add_argument(
+            '--server_only', action='store_true',
+            help='Limit filtering to only server tests.')
+
+    action_group = parser.add_argument_group(
+            'actions', 'Actions which are performed on a control file.')
+    action_group.add_argument(
+            '--remove_contacts', required=False, action='append',
+            metavar='EMAIL,EMAIL,...',
+            help=('Remove the given (comma or newline-separated) '
+                  'email addresses from Contacts.'))
+    action_group.add_argument(
+            '--append_contacts', required=False, action='append',
+            metavar='EMAIL',
+            help=('Add (or move) the given email addresses to the END of the '
+                  'contacts list.'))
+    action_group.add_argument(
+            '--prepend_contacts', required=False, action='append',
+            metavar='EMAIL',
+            help=('Add (or move) the given email addresses to the START of the '
+                  'contacts list.'))
+
+    filter_group = parser.add_argument_group(
+            'filters', 'Filters to specify which control files are touched.')
+    filter_group.add_argument(
+            '--test_names', required=False, action='append',
+            help=('Action: Modify only the given (comma or newline-separated) '
+                  'test ids (i.e. inlcuding \'tauto.\' prefix).'))
+
+    return parser.parse_args()
+
+def _split_list_input(string_input):
+    """Split the given string into comma or newline-separated values."""
+    return [elt.strip() for elt in re.split(r',\s*|\n\s*', string_input)]
+
+def _get_actions(args):
+    """Given the arguments to the CLI, return a list of actions requested."""
+    output = []
+    if args.remove_contacts:
+        for elt in args.remove_contacts:
+            output.append(actions.remove_contacts(_split_list_input(elt)))
+    if args.append_contacts:
+        for elt in args.append_contact:
+            output.append(actions.append_contacts(_split_list_input(elt)))
+    if args.prepend_contacts:
+        for elt in args.prepend_contact:
+            output.append(actions.prepend_contacts(_split_list_input(elt)))
+    return output
+
+
+def _get_filters(args):
+    """Given the arguments to the CLI, return a list of filters requested."""
+    output = []
+    if args.test_names:
+        for elt in args.test_names:
+            tests = [t.strip() for t in re.split(r',\s*|\n\s*', elt)]
+            output.append(filters.test_list(tests))
+    return output
+
+
 def main():
-    os.chdir(SRC_DIR)
-    action_func = actions.remove_contacts(['notarealemail@google.com'])
-    filter_func = filters.all_tests()
-    dry_run = True
-    modify_control_files([action_func], [filter_func], dry_run)
+    args = _set_up_args()
+    print(args)
+
+    modify_control_files(
+            args.src_dir, _get_actions(args), _get_filters(args),
+            dry_run=not args.write_out,
+            public=not args.private_only, private=not args.public_only,
+            client=not args.server_only, server=not args.client_only)
 
 if __name__ == '__main__':
     main()
diff --git a/metadata_modifier/tauto_modify/test_actions.py b/metadata_modifier/tauto_modify/test_actions.py
index 4afc08f..d9c2671 100644
--- a/metadata_modifier/tauto_modify/test_actions.py
+++ b/metadata_modifier/tauto_modify/test_actions.py
@@ -22,13 +22,48 @@
 
         self.assertTrue('contacts' in cf.metadata)
         self.assertTrue(delete_me in cf.metadata['contacts'])
-        modified = action(cf)
-        self.assertTrue(modified)
+        self.assertTrue(action(cf))
         self.assertTrue('contacts' in cf.metadata)
         self.assertFalse(delete_me in cf.metadata['contacts'])
 
-        modified = action(cf)
-        self.assertFalse(modified)
+        self.assertFalse(action(cf))
+
+    def test_append_contacts(self):
+        append_1 = "appendable1@google.com"
+        append_2 = "appendable2@google.com"
+        action = actions.append_contacts([append_1, append_2])
+        test_file = os.path.join(TEST_DATA_DIR, 'control.actions')
+        cf = cf_parse.ControlFile(test_file)
+        starting_len = len(cf.metadata['contacts'])
+
+        self.assertTrue(action(cf))
+        self.assertTrue('contacts' in cf.metadata)
+        self.assertTrue(append_1 in cf.metadata['contacts'])
+        self.assertTrue(append_2 in cf.metadata['contacts'])
+        self.assertEqual(cf.metadata['contacts'].index(append_1), starting_len)
+        self.assertEqual(cf.metadata['contacts'].index(append_2),
+                         starting_len + 1)
+        self.assertEqual(len(cf.metadata['contacts']), starting_len + 2)
+
+        self.assertTrue(action(cf))
+        self.assertEqual(len(cf.metadata['contacts']), starting_len + 2)
+
+    def test_prepend_contacts(self):
+        prepend = "prependable@google.com"
+        action = actions.prepend_contacts([prepend])
+        test_file = os.path.join(TEST_DATA_DIR, 'control.actions')
+        cf = cf_parse.ControlFile(test_file)
+        starting_len = len(cf.metadata['contacts'])
+
+        self.assertTrue(action(cf))
+        self.assertTrue('contacts' in cf.metadata)
+        self.assertTrue(prepend in cf.metadata['contacts'])
+        self.assertEqual(cf.metadata['contacts'].index(prepend), 0)
+        self.assertEqual(len(cf.metadata['contacts']), starting_len + 1)
+
+        self.assertTrue(action(cf))
+        self.assertEqual(len(cf.metadata['contacts']), starting_len + 1)
+
 
 if __name__ == '__main__':
     unittest.main()