[autotest] Manage shards with atest

This adds functionality to list, create and delete shards to atest.

BUG=None
TEST=Ran suites, manual test of create/list/delete.

Change-Id: I0771d0c1b46c7c6890819822204f42ac2211b104
Reviewed-on: https://chromium-review.googlesource.com/218295
Tested-by: Jakob Jülich <[email protected]>
Reviewed-by: Prashanth B <[email protected]>
Commit-Queue: Jakob Jülich <[email protected]>
diff --git a/cli/atest_unittest.py b/cli/atest_unittest.py
index b4f62d6..29a8c75 100755
--- a/cli/atest_unittest.py
+++ b/cli/atest_unittest.py
@@ -26,8 +26,9 @@
     def test_main_help(self):
         """Main help level"""
         self._test_help(argv=['atest'],
-                        out_words_ok=['atest [acl|host|job|label|atomicgroup'
-                                      '|test|user] [action] [options]'],
+                        out_words_ok=['atest [acl|host|job|label|shard'
+                                      '|atomicgroup|test|user] '
+                                      '[action] [options]'],
                         err_words_ok=[])
 
 
@@ -50,14 +51,14 @@
     def test_main_no_topic(self):
         self.run_cmd(['atest'], exit_code=1,
                      out_words_ok=['atest '
-                                   '[acl|host|job|label|atomicgroup|test|user] '
-                                   '[action] [options]'],
+                                   '[acl|host|job|label|shard|atomicgroup|test'
+                                   '|user] [action] [options]'],
                      err_words_ok=['No topic argument'])
 
 
     def test_main_bad_topic(self):
         self.run_cmd(['atest', 'bad_topic'], exit_code=1,
-                     out_words_ok=['atest [acl|host|job|label|atomicgroup'
+                     out_words_ok=['atest [acl|host|job|label|shard|atomicgroup'
                                    '|test|user] [action] [options]'],
                      err_words_ok=['Invalid topic bad_topic\n'])
 
diff --git a/cli/shard.py b/cli/shard.py
new file mode 100644
index 0000000..1f7addb
--- /dev/null
+++ b/cli/shard.py
@@ -0,0 +1,140 @@
+#
+# Copyright 2008 Google Inc. All Rights Reserved.
+
+"""
+The shard module contains the objects and methods used to
+manage shards in Autotest.
+
+The valid actions are:
+create:      creates shard
+remove:      deletes shard(s)
+list:        lists shards with label
+
+See topic_common.py for a High Level Design and Algorithm.
+"""
+
+import os, sys
+from autotest_lib.cli import topic_common, action_common
+
+
+class shard(topic_common.atest):
+    """shard class
+    atest shard [create|delete|list] <options>"""
+    usage_action = '[create|delete|list]'
+    topic = msg_topic = 'shard'
+    msg_items = '<shards>'
+
+    def __init__(self):
+        """Add to the parser the options common to all the
+        shard actions"""
+        super(shard, self).__init__()
+
+        self.topic_parse_info = topic_common.item_parse_info(
+            attribute_name='shards',
+            use_leftover=True)
+
+
+    def get_items(self):
+        return self.shards
+
+
+class shard_help(shard):
+    """Just here to get the atest logic working.
+    Usage is set by its parent"""
+    pass
+
+
+class shard_list(action_common.atest_list, shard):
+    """Class for running atest shard list [--label <labels>]"""
+
+    def parse(self):
+        host_info = topic_common.item_parse_info(attribute_name='labels',
+                                                 inline_option='labels')
+        return super(shard_list, self).parse([host_info])
+
+
+    def execute(self):
+        return super(shard_list, self).execute(op='get_shards')
+
+
+    def warn_if_label_assigned_to_multiple_shards(self, results):
+        """Prints a warning if one label is assigned to multiple shards.
+
+        This should never happen, but if it does, better be safe.
+
+        @param results: Results as passed to output().
+        """
+        assigned_labels = set()
+        for line in results:
+            for label in line['labels']:
+                if label in assigned_labels:
+                    sys.stderr.write('WARNING: label %s is assigned to '
+                                     'multiple shards.\n'
+                                     'This will lead to unpredictable behavor '
+                                     'in which hosts and jobs will be assigned '
+                                     'to which shard.\n')
+                assigned_labels.add(label)
+
+
+    def output(self, results):
+        self.warn_if_label_assigned_to_multiple_shards(results)
+        super(shard_list, self).output(results, ['hostname', 'labels'])
+
+
+class shard_create(action_common.atest_create, shard):
+    """Class for running atest shard create -l <label> <shard>"""
+    def __init__(self):
+        super(shard_create, self).__init__()
+        self.parser.add_option('-l', '--label',
+                               help=('Assign LABEL to the SHARD. All jobs that '
+                                     'require this label, will be run on the '
+                                     'shard.'),
+                               type='string',
+                               metavar='LABEL')
+
+
+    def parse(self):
+        (options, leftover) = super(shard_create,
+                                    self).parse(req_items='shards')
+        if not options.label:
+            print 'Must provide a label with -l <label>'
+            self.parser.print_help()
+            sys.exit(1)
+        self.data_item_key = 'hostname'
+        self.data['label'] = options.label
+        return (options, leftover)
+
+
+class shard_delete(action_common.atest_delete, shard):
+    """Class for running atest shard delete <shards>"""
+    def __init__(self):
+        super(shard_delete, self).__init__()
+        self.parser.add_option('-y', '--yes',
+                               help=('Answer all questions with yes.'),
+                               action='store_true',
+                               metavar='LABEL')
+
+
+    def parse(self):
+        (options, leftover) = super(shard_delete, self).parse()
+        self.yes = options.yes
+        self.data_item_key = 'hostname'
+        return (options, leftover)
+
+
+    def execute(self, *args, **kwargs):
+        if self.yes or self._prompt_confirmation():
+            return super(shard_delete, self).execute(*args, **kwargs)
+        print 'Aborting.'
+        return []
+
+
+    def _prompt_confirmation(self):
+        print 'Please ensure the shard host is powered off.'
+        print ('Otherwise DUTs might be used by multiple shards at the same '
+               'time, which will lead to serious correctness problems.')
+        sys.stdout.write('Continue? [y/N] ')
+        read = raw_input().lower()
+        if read == 'y':
+            return True
+        return False
diff --git a/cli/shard_unittest.py b/cli/shard_unittest.py
new file mode 100755
index 0000000..a72992a
--- /dev/null
+++ b/cli/shard_unittest.py
@@ -0,0 +1,61 @@
+#!/usr/bin/python
+#
+# Copyright 2008 Google Inc. All Rights Reserved.
+
+"""Tests for shard."""
+
+import unittest, sys, os
+
+import common
+from autotest_lib.cli import cli_mock, topic_common
+
+
+class shard_list_unittest(cli_mock.cli_unittest):
+    values = [{'hostname': u'shard1', u'id': 1, 'labels': ['board:lumpy']},
+              {'hostname': u'shard2', u'id': 3, 'labels': ['board:daisy']},
+              {'hostname': u'shard3', u'id': 5, 'labels': ['board:stumpy']},
+              {'hostname': u'shard4', u'id': 6, 'labels': ['board:link']}]
+
+
+    def test_shard_list(self):
+        self.run_cmd(argv=['atest', 'shard', 'list'],
+                     rpcs=[('get_shards', {}, True, self.values)],
+                     out_words_ok=['shard1', 'shard2', 'shard3', 'shard4'],
+                     out_words_no=['plat0', 'plat1'])
+
+
+class shard_create_unittest(cli_mock.cli_unittest):
+    def test_execute_create_two_shards(self):
+        self.run_cmd(argv=['atest', 'shard', 'create',
+                           '-l', 'board:lumpy', 'shard0'],
+                     rpcs=[('add_shard',
+                            {'hostname': 'shard0', 'label': 'board:lumpy'},
+                            True, 42)],
+                     out_words_ok=['Created', 'shard0'])
+
+
+    def test_execute_create_two_shards_bad(self):
+        self.run_cmd(argv=['atest', 'shard', 'create',
+                           '-l', 'board:lumpy', 'shard0'],
+                     rpcs=[('add_shard',
+                            {'hostname': 'shard0', 'label': 'board:lumpy'},
+                            False,
+                            '''ValidationError: {'name':
+                            'This value must be unique (shard1)'}''')],
+                     out_words_no=['shard0'],
+                     err_words_ok=['shard0', 'ValidationError'])
+
+
+class shard_delete_unittest(cli_mock.cli_unittest):
+    def test_execute_delete_shards(self):
+        self.run_cmd(argv=['atest', 'shard', 'delete',
+                           'shard0', 'shard1', '--yes'],
+                     rpcs=[('delete_shard', {'hostname': 'shard0'}, True, None),
+                           ('delete_shard', {'hostname': 'shard1'}, True, None)
+                           ],
+                     out_words_ok=['Deleted', 'shard0', 'shard1'])
+
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/cli/topic_common.py b/cli/topic_common.py
index d2fd891..c44f67d 100644
--- a/cli/topic_common.py
+++ b/cli/topic_common.py
@@ -162,7 +162,7 @@
                  filename_option='', use_leftover=False):
         """Object keeping track of the parsing options that will
         make up the content of the atest attribute:
-        atttribute_name: the atest attribute name to populate    (label)
+        attribute_name: the atest attribute name to populate    (label)
         inline_option: the option containing the items           (--label)
         filename_option: the option containing the filename      (--blist)
         use_leftover: whether to add the leftover arguments or not."""
@@ -241,7 +241,7 @@
     Should only be instantiated by itself for usage
     references, otherwise, the <topic> objects should
     be used."""
-    msg_topic = "[acl|host|job|label|atomicgroup|test|user]"
+    msg_topic = "[acl|host|job|label|shard|atomicgroup|test|user]"
     usage_action = "[action]"
     msg_items = ''