Cleaned up host mod and host create logic.

Merged site_host.py into host.py. Created base class for shared logic
between mod and create actions. Also merged tests to run the same test
cases against both actions.

CQ-DEPEND=CL:354552

BUG=None
TEST=Unit tests verified locally.

Change-Id: If14786ef5f58a415406efd7a2f24b96ca3ea1a74
Reviewed-on: https://chromium-review.googlesource.com/355733
Commit-Ready: Justin Giorgi <jgiorgi@google.com>
Tested-by: Justin Giorgi <jgiorgi@chromium.org>
Reviewed-by: Simran Basi <sbasi@chromium.org>
diff --git a/cli/host.py b/cli/host.py
index 26d4d76..52736f3 100644
--- a/cli/host.py
+++ b/cli/host.py
@@ -1,4 +1,3 @@
-#
 # Copyright 2008 Google Inc. All Rights Reserved.
 
 """
@@ -20,10 +19,14 @@
 See topic_common.py for a High Level Design and Algorithm.
 
 """
+import common
 import re
+import socket
 
-from autotest_lib.cli import action_common, topic_common
-from autotest_lib.client.common_lib import host_protections
+from autotest_lib.cli import action_common, rpc, topic_common
+from autotest_lib.client.bin import utils
+from autotest_lib.client.common_lib import error, host_protections
+from autotest_lib.server import hosts
 
 
 class host(topic_common.atest):
@@ -310,29 +313,21 @@
                                                 'job_name',
                                                 'status'])
 
-
-class host_mod(host):
-    """atest host mod --lock|--unlock|--force_modify_locking|--protection
-    --mlist <file>|<hosts>"""
-    usage_action = 'mod'
+class BaseHostModCreate(host):
     attribute_regex = r'^(?P<attribute>\w+)=(?P<value>.+)?'
+    attr_split_regex = r'[^\\],' # Matches , not preceeded by \
 
     def __init__(self):
-        """Add the options specific to the mod action"""
-        self.data = {}
+        """Add the options shared between host mod and host create actions."""
         self.messages = []
-        self.attribute = None
-        self.value = None
-        super(host_mod, self).__init__()
+        self.host_ids = {}
+        super(BaseHostModCreate, self).__init__()
         self.parser.add_option('-l', '--lock',
                                help='Lock hosts',
                                action='store_true')
         self.parser.add_option('-u', '--unlock',
                                help='Unlock hosts',
                                action='store_true')
-        self.parser.add_option('-f', '--force_modify_locking',
-                               help='Forcefully lock\unlock a host',
-                               action='store_true')
         self.parser.add_option('-r', '--lock_reason',
                                help='Reason for locking hosts',
                                default='')
@@ -342,87 +337,11 @@
                                      ', '.join('"%s"' % p
                                                for p in self.protections)),
                                choices=self.protections)
-        self.parser.add_option('--attribute', '-a', default='',
-                               help=('Host attribute to add or change. Format '
-                                     'is <attribute>=<value>. Value can be '
-                                     'blank to delete attribute.'))
-
-
-    def parse(self):
-        """Consume the specific options"""
-        (options, leftover) = super(host_mod, self).parse()
-
-        self._parse_lock_options(options)
-        if options.force_modify_locking:
-             self.data['force_modify_locking'] = True
-
-        if options.protection:
-            self.data['protection'] = options.protection
-            self.messages.append('Protection set to "%s"' % options.protection)
-
-        if len(self.data) == 0 and not options.attribute:
-            self.invalid_syntax('No modification requested')
-
-        if options.attribute:
-            match = re.match(self.attribute_regex, options.attribute)
-            if not match:
-                self.invalid_syntax('Attributes must be in <attribute>=<value>'
-                                    ' syntax!')
-
-            self.attribute = match.group('attribute')
-            self.value = match.group('value')
-
-        return (options, leftover)
-
-
-    def execute(self):
-        successes = []
-        for host in self.hosts:
-            try:
-                res = self.execute_rpc('modify_host', item=host,
-                                       id=host, **self.data)
-                if self.attribute:
-                    self.execute_rpc('set_host_attribute',
-                                     attribute=self.attribute,
-                                     value=self.value, hostname=host)
-                # TODO: Make the AFE return True or False,
-                # especially for lock
-                successes.append(host)
-            except topic_common.CliError, full_error:
-                # Already logged by execute_rpc()
-                pass
-
-        return successes
-
-
-    def output(self, hosts):
-        for msg in self.messages:
-            self.print_wrapped(msg, hosts)
-
-
-class host_create(host):
-    """atest host create [--lock|--unlock --platform <arch>
-    --labels <labels>|--blist <label_file>
-    --acls <acls>|--alist <acl_file>
-    --protection <protection_type>
-    --mlist <mach_file>] <hosts>"""
-    usage_action = 'create'
-
-    def __init__(self):
-        self.messages = []
-        super(host_create, self).__init__()
-        self.parser.add_option('-l', '--lock',
-                               help='Create the hosts as locked',
-                               action='store_true', default=False)
-        self.parser.add_option('-u', '--unlock',
-                               help='Create the hosts as '
-                               'unlocked (default)',
-                               action='store_true')
-        self.parser.add_option('-r', '--lock_reason',
-                               help='Reason for locking hosts',
-                               default='')
-        self.parser.add_option('-t', '--platform',
-                               help='Sets the platform label')
+        self.parser.add_option('--attributes', '-i', default='',
+                               help=('Host attributes to add or change. Format '
+                                     'is <attribute>=<value>. Comma delimited '
+                                     'for multiple attributes. Use \\\\ (two '
+                                     'backslashes) to escape delimitter.'))
         self.parser.add_option('-b', '--labels',
                                help='Comma separated list of labels')
         self.parser.add_option('-B', '--blist',
@@ -435,18 +354,13 @@
                                help='File listing the acls',
                                type='string',
                                metavar='ACL_FLIST')
-        self.parser.add_option('-p', '--protection', type='choice',
-                               help=('Set the protection level on a host.  '
-                                     'Must be one of: %s' %
-                                     ', '.join('"%s"' % p
-                                               for p in self.protections)),
-                               choices=self.protections)
-        self.parser.add_option('-s', '--serials',
-                               help=('Comma separated list of adb-based device '
-                                     'serials'))
+        self.parser.add_option('-t', '--platform',
+                               help='Sets the platform label')
 
 
     def parse(self):
+        """Consume the options common to host create and host mod.
+        """
         label_info = topic_common.item_parse_info(attribute_name='labels',
                                                  inline_option='labels',
                                                  filename_option='blist')
@@ -454,51 +368,312 @@
                                                 inline_option='acls',
                                                 filename_option='alist')
 
-        (options, leftover) = super(host_create, self).parse([label_info,
+        (options, leftover) = super(BaseHostModCreate, self).parse([label_info,
                                                               acl_info],
                                                              req_items='hosts')
 
         self._parse_lock_options(options)
-        self.locked = options.lock
-        self.platform = getattr(options, 'platform', None)
-        self.serials = getattr(options, 'serials', None)
-        if self.serials:
-            if len(self.hosts) > 1:
-                raise topic_common.CliError('Can not specify serials with '
-                                            'multiple hosts')
-            self.serials = self.serials.split(',')
+
         if options.protection:
             self.data['protection'] = options.protection
+            self.messages.append('Protection set to "%s"' % options.protection)
+
+        self.attributes = {}
+        if options.attributes:
+            # Extract pairs of attributes
+            last_end = 0
+            groups = []
+            for m in re.finditer(self.attr_split_regex, options.attributes):
+                # The first char of the match must be included because it is
+                # the char before the delimitter
+                groups.append(options.attributes[last_end:m.start()+1])
+                last_end = m.end()
+            if options.attributes[last_end:]:
+                groups.append(options.attributes[last_end:])
+
+            # Process pairs of attributes
+            for group in groups:
+                match = re.match(self.attribute_regex, group)
+                if not match:
+                    self.invalid_syntax('Attributes must be in '
+                                        '<attribute>=<value> syntax!')
+
+                attribute = match.group('attribute')
+                value = match.group('value')
+
+                if attribute in self.attributes:
+                    raise topic_common.CliError('Multiple values provided for '
+                                                'attribute %s.' % attribute)
+                self.attributes[attribute] = value
+
+        self.platform = options.platform
         return (options, leftover)
 
 
+    def _set_acls(self, hosts, acls):
+        """Add hosts to acls (and remove from all other acls).
+
+        @param hosts: list of hostnames
+        @param acls: list of acl names
+        """
+        # Remove from all ACLs except 'Everyone' and ACLs in list
+        # Skip hosts that don't exist
+        for host in hosts:
+            if host not in self.host_ids:
+                continue
+            host_id = self.host_ids[host]
+            for a in self.execute_rpc('get_acl_groups', hosts=host_id):
+                if a['name'] not in self.acls and a['id'] != 1:
+                    self.execute_rpc('acl_group_remove_hosts', id=a['id'],
+                                     hosts=self.hosts)
+
+        # Add hosts to the ACLs
+        self.check_and_create_items('get_acl_groups', 'add_acl_group',
+                                    self.acls)
+        for a in acls:
+            self.execute_rpc('acl_group_add_hosts', id=a, hosts=hosts)
+
+
+    def _remove_labels(self, host, condition):
+        """Remove all labels from host that meet condition(label).
+
+        @param host: hostname
+        @param condition: callable that returns bool when given a label
+        """
+        if host in self.host_ids:
+            host_id = self.host_ids[host]
+            labels_to_remove = []
+            for l in self.execute_rpc('get_labels', host=host_id):
+                if condition(l):
+                    labels_to_remove.append(l['id'])
+            if labels_to_remove:
+                self.execute_rpc('host_remove_labels', id=host_id,
+                                 labels=labels_to_remove)
+
+
+    def _set_labels(self, host, labels):
+        """Apply labels to host (and remove all other labels).
+
+        @param host: hostname
+        @param labels: list of label names
+        """
+        condition = lambda l: l['name'] not in labels and not l['platform']
+        self._remove_labels(host, condition)
+        self.check_and_create_items('get_labels', 'add_label', labels)
+        self.execute_rpc('host_add_labels', id=host, labels=labels)
+
+
+    def _set_platform_label(self, host, platform_label):
+        """Apply the platform label to host (and remove existing).
+
+        @param host: hostname
+        @param platform_label: platform label's name
+        """
+        self._remove_labels(host, lambda l: l['platform'])
+        self.check_and_create_items('get_labels', 'add_label', [platform_label],
+                                    platform=True)
+        self.execute_rpc('host_add_labels', id=host, labels=[platform_label])
+
+
+    def _set_attributes(self, host, attributes):
+        """Set attributes on host.
+
+        @param host: hostname
+        @param attributes: attribute dictionary
+        """
+        for attr, value in self.attributes.iteritems():
+            self.execute_rpc('set_host_attribute', attribute=attr,
+                             value=value, hostname=host)
+
+
+class host_mod(BaseHostModCreate):
+    """atest host mod [--lock|--unlock --force_modify_locking
+    --platform <arch>
+    --labels <labels>|--blist <label_file>
+    --acls <acls>|--alist <acl_file>
+    --protection <protection_type>
+    --attributes <attr>=<value>;<attr>=<value>
+    --mlist <mach_file>] <hosts>"""
+    usage_action = 'mod'
+    attribute_regex = r'^(?P<attribute>\w+)=(?P<value>.+)?'
+
+    def __init__(self):
+        """Add the options specific to the mod action"""
+        super(host_mod, self).__init__()
+        self.parser.add_option('-f', '--force_modify_locking',
+                               help='Forcefully lock\unlock a host',
+                               action='store_true')
+        self.parser.add_option('--remove_acls',
+                               help='Remove all active acls.',
+                               action='store_true')
+        self.parser.add_option('--remove_labels',
+                               help='Remove all labels.',
+                               action='store_true')
+
+
+    def parse(self):
+        """Consume the specific options"""
+        (options, leftover) = super(host_mod, self).parse()
+
+        if options.force_modify_locking:
+             self.data['force_modify_locking'] = True
+
+        self.remove_acls = options.remove_acls
+        self.remove_labels = options.remove_labels
+
+        return (options, leftover)
+
+
+    def execute(self):
+        successes = []
+        for host in self.execute_rpc('get_hosts', hostname__in=self.hosts):
+            self.host_ids[host['hostname']] = host['id']
+        for host in self.hosts:
+            if host not in self.host_ids:
+                self.failure('Cannot modify non-existant host %s.' % host)
+                continue
+            host_id = self.host_ids[host]
+
+            try:
+                if self.data:
+                    self.execute_rpc('modify_host', item=host,
+                                     id=host, **self.data)
+
+                if self.attributes:
+                    self._set_attributes(host, self.attributes)
+
+                if self.labels or self.remove_labels:
+                    self._set_labels(host, self.labels)
+
+                if self.platform:
+                    self._set_platform_label(host, self.platform)
+
+                # TODO: Make the AFE return True or False,
+                # especially for lock
+                successes.append(host)
+            except topic_common.CliError, full_error:
+                # Already logged by execute_rpc()
+                pass
+
+        if self.acls or self.remove_acls:
+            self._set_acls(self.hosts, self.acls)
+
+        return successes
+
+
+    def output(self, hosts):
+        for msg in self.messages:
+            self.print_wrapped(msg, hosts)
+
+
+class HostInfo(object):
+    """Store host information so we don't have to keep looking it up."""
+    def __init__(self, hostname, platform, labels):
+        self.hostname = hostname
+        self.platform = platform
+        self.labels = labels
+
+
+class host_create(BaseHostModCreate):
+    """atest host create [--lock|--unlock --platform <arch>
+    --labels <labels>|--blist <label_file>
+    --acls <acls>|--alist <acl_file>
+    --protection <protection_type>
+    --attributes <attr>=<value>;<attr>=<value>
+    --mlist <mach_file>] <hosts>"""
+    usage_action = 'create'
+
+    def parse(self):
+        """Option logic specific to create action.
+        """
+        (options, leftovers) = super(host_create, self).parse()
+        self.locked = options.lock
+        if 'serials' in self.attributes:
+            if len(self.hosts) > 1:
+                raise topic_common.CliError('Can not specify serials with '
+                                            'multiple hosts.')
+
+
+    @classmethod
+    def construct_without_parse(
+            cls, web_server, hosts, platform=None,
+            locked=False, lock_reason='', labels=[], acls=[],
+            protection=host_protections.Protection.NO_PROTECTION):
+        """Construct a host_create object and fill in data from args.
+
+        Do not need to call parse after the construction.
+
+        Return an object of site_host_create ready to execute.
+
+        @param web_server: A string specifies the autotest webserver url.
+            It is needed to setup comm to make rpc.
+        @param hosts: A list of hostnames as strings.
+        @param platform: A string or None.
+        @param locked: A boolean.
+        @param lock_reason: A string.
+        @param labels: A list of labels as strings.
+        @param acls: A list of acls as strings.
+        @param protection: An enum defined in host_protections.
+        """
+        obj = cls()
+        obj.web_server = web_server
+        try:
+            # Setup stuff needed for afe comm.
+            obj.afe = rpc.afe_comm(web_server)
+        except rpc.AuthError, s:
+            obj.failure(str(s), fatal=True)
+        obj.hosts = hosts
+        obj.platform = platform
+        obj.locked = locked
+        if locked and lock_reason.strip():
+            obj.data['lock_reason'] = lock_reason.strip()
+        obj.labels = labels
+        obj.acls = acls
+        if protection:
+            obj.data['protection'] = protection
+        obj.attributes = {}
+        return obj
+
+
     def _execute_add_one_host(self, host):
         # Always add the hosts as locked to avoid the host
         # being picked up by the scheduler before it's ACL'ed.
-        # We enforce lock reasons for each lock, so we
-        # provide a 'dummy' if we are intending to unlock after.
         self.data['locked'] = True
         if not self.locked:
             self.data['lock_reason'] = 'Forced lock on device creation'
-        self.execute_rpc('add_host', hostname=host,
-                         status="Ready", **self.data)
+        self.execute_rpc('add_host', hostname=host, status="Ready", **self.data)
 
-        # Now add the platform label
-        labels = self.labels[:]
-        if self.platform:
-            labels.append(self.platform)
-        if len (labels):
-            self.execute_rpc('host_add_labels', id=host, labels=labels)
+        # If there are labels avaliable for host, use them.
+        host_info = self.host_info_map[host]
+        labels = set(self.labels)
+        if host_info.labels:
+            labels.update(host_info.labels)
+
+        if labels:
+            self._set_labels(host, list(labels))
+
+        # Now add the platform label.
+        # If a platform was not provided and we were able to retrieve it
+        # from the host, use the retrieved platform.
+        platform = self.platform if self.platform else host_info.platform
+        if platform:
+            self._set_platform_label(host, platform)
+
+        if self.attributes:
+            self._set_attributes(host, self.attributes)
 
 
     def _execute_add_hosts(self):
-        successful_hosts = self.site_create_hosts_hook()
+        successful_hosts = []
+        for host in self.hosts:
+            try:
+                self._execute_add_one_host(host)
+                successful_hosts.append(host)
+            except topic_common.CliError:
+                pass
 
         if successful_hosts:
-            for acl in self.acls:
-                self.execute_rpc('acl_group_add_hosts',
-                                 id=acl,
-                                 hosts=successful_hosts)
+            self._set_acls(successful_hosts, self.acls)
 
             if not self.locked:
                 for host in successful_hosts:
@@ -508,36 +683,33 @@
 
 
     def execute(self):
-        # We need to check if these labels & ACLs exist,
-        # and create them if not.
-        if self.platform:
-            self.check_and_create_items('get_labels', 'add_label',
-                                        [self.platform],
-                                        platform=True)
-
-        if self.labels:
-            self.check_and_create_items('get_labels', 'add_label',
-                                        self.labels,
-                                        platform=False)
-
-        if self.acls:
-            self.check_and_create_items('get_acl_groups',
-                                        'add_acl_group',
-                                        self.acls)
-
-        return self._execute_add_hosts()
-
-
-    def site_create_hosts_hook(self):
-        successful_hosts = []
+        # Check to see if the platform or any other labels can be grabbed from
+        # the hosts.
+        self.host_info_map = {}
         for host in self.hosts:
             try:
-                self._execute_add_one_host(host)
-                successful_hosts.append(host)
-            except topic_common.CliError:
-                pass
+                if utils.ping(host, tries=1, deadline=1) == 0:
+                    serials = self.attributes.get('serials', '').split(',')
+                    if serials and len(serials) > 1:
+                        host_dut = hosts.create_testbed(host,
+                                                        adb_serials=serials)
+                    else:
+                        adb_serial = self.attributes.get('serials')
+                        host_dut = hosts.create_host(host,
+                                                     adb_serial=adb_serial)
+                    host_info = HostInfo(host, host_dut.get_platform(),
+                                         host_dut.get_labels())
+                else:
+                    # Can't ping the host, use default information.
+                    host_info = HostInfo(host, None, [])
+            except (socket.gaierror, error.AutoservRunError,
+                    error.AutoservSSHTimeout):
+                # We may be adding a host that does not exist yet or we can't
+                # reach due to hostname/address issues or if the host is down.
+                host_info = HostInfo(host, None, [])
+            self.host_info_map[host] = host_info
 
-        return successful_hosts
+        return self._execute_add_hosts()
 
 
     def output(self, hosts):
diff --git a/cli/host_unittest.py b/cli/host_unittest.py
index beda185..92fb560 100755
--- a/cli/host_unittest.py
+++ b/cli/host_unittest.py
@@ -1272,165 +1272,354 @@
                                    'testjob'])
 
 
-class host_mod_unittest(cli_mock.cli_unittest):
-    def test_execute_lock_one_host(self):
-        self.run_cmd(argv=['atest', 'host', 'mod', '--lock', 'host0'],
-                     rpcs=[('modify_host', {'id': 'host0', 'locked': True},
-                            True, None)],
-                     out_words_ok=['Locked', 'host0'])
+class host_mod_create_tests(object):
+
+    def _gen_attributes_rpcs(self, host, attributes):
+        """Generate RPCs expected to add attributes to host.
+
+        @param host: hostname
+        @param attributes: dict of attributes
+
+        @return: list of rpcs to expect
+        """
+        rpcs = []
+        for attr, val in attributes.iteritems():
+            rpcs.append(('set_host_attribute',
+                         {
+                             'hostname': host,
+                             'attribute': attr,
+                             'value': val,
+                         },
+                         True, None))
+        return rpcs
 
 
-    def test_execute_unlock_two_hosts(self):
-        self.run_cmd(argv=['atest', 'host', 'mod', '-u', 'host0,host1'],
-                     rpcs=[('modify_host', {'id': 'host1', 'locked': False,
-                                            'lock_reason': ''},
-                            True, None),
-                           ('modify_host', {'id': 'host0', 'locked': False,
-                                            'lock_reason': ''},
-                            True, None)],
-                     out_words_ok=['Unlocked', 'host0', 'host1'])
+    def _gen_labels_rpcs(self, labels, platform=False, host_id=None):
+        """Generate RPCS expected to add labels.
+
+        @param labels: list of label names
+        @param platform: labels are platform labels
+        @param host_id: Host id old labels will be deleted from (if host exists)
+        """
+        rpcs = []
+        if host_id:
+            rpcs.append(('get_labels', {'host': host_id}, True, []))
+        for label in labels:
+            rpcs += [
+                ('get_labels', {'name': label}, True, []),
+                ('add_label', {'name': label}, True, None)
+            ]
+            if platform:
+                rpcs[-1][1]['platform'] = True
+        return rpcs
 
 
-    def test_execute_force_lock_one_host(self):
-        self.run_cmd(argv=['atest', 'host', 'mod', '--lock',
-                           '--force_modify_locking', 'host0'],
-                     rpcs=[('modify_host',
-                            {'id': 'host0', 'locked': True,
-                             'force_modify_locking': True},
-                            True, None)],
-                     out_words_ok=['Locked', 'host0'])
+    def _gen_acls_rpcs(self, hosts, acls, host_ids=[]):
+        """Generate RPCs expected to add acls.
+
+        @param hosts: list of hostnames
+        @param acls: list of acl names
+        @param host_ids: List of host_ids if hosts already exist
+        """
+        rpcs = []
+        for host_id in host_ids:
+            rpcs.append(('get_acl_groups', {'hosts': host_id}, True, []))
+        for acl in acls:
+            rpcs.append(('get_acl_groups', {'name': acl}, True, []))
+            rpcs.append(('add_acl_group', {'name': acl}, True, None))
+        for acl in acls:
+            rpcs.append((
+                'acl_group_add_hosts',
+                {
+                    'hosts': hosts,
+                    'id': acl,
+                },
+                True,
+                None,
+            ))
+        return rpcs
 
 
-    def test_execute_force_unlock_one_host(self):
-        self.run_cmd(argv=['atest', 'host', 'mod', '--unlock',
-                           '--force_modify_locking', 'host0'],
-                     rpcs=[('modify_host',
-                            {'id': 'host0', 'locked': False,
-                             'force_modify_locking': True,
-                             'lock_reason': ''},
-                            True, None)],
-                     out_words_ok=['Unlocked', 'host0'])
+    def test_lock_one_host(self):
+        """Test locking host / creating host locked."""
+        lock_reason = 'Because'
+        rpcs, out = self._gen_expectations(locked=True, lock_reason=lock_reason)
+        self.run_cmd(argv=self._command_single + ['--lock', '--lock_reason',
+                                                  lock_reason],
+                     rpcs=rpcs, out_words_ok=out)
 
 
-    def test_execute_lock_unknown_hosts(self):
-        self.run_cmd(argv=['atest', 'host', 'mod', '-l', 'host0,host1',
-                           'host2'],
-                     rpcs=[('modify_host', {'id': 'host2', 'locked': True},
-                            True, None),
-                           ('modify_host', {'id': 'host1', 'locked': True},
-                            False, 'DoesNotExist: Host matching '
-                            'query does not exist.'),
-                           ('modify_host', {'id': 'host0', 'locked': True},
-                            True, None)],
-                     out_words_ok=['Locked', 'host0', 'host2'],
-                     err_words_ok=['Host', 'matching', 'query', 'host1'])
+    def test_unlock_multiple_hosts(self):
+        """Test unlocking host / creating host unlocked."""
+        rpcs, out = self._gen_expectations(hosts=self._hosts, locked=False)
+        self.run_cmd(argv=self._command_multiple + ['--unlock'], rpcs=rpcs,
+                     out_words_ok=out)
 
 
-    def test_execute_protection_hosts(self):
-        mfile = cli_mock.create_file('host0\nhost1,host2\nhost3 host4')
+    def test_machine_list(self):
+        """Test action an machines from machine list file."""
+        mfile = cli_mock.create_file(','.join(self._hosts))
+        rpcs, out = self._gen_expectations(hosts=self._hosts, locked=False)
         try:
-            self.run_cmd(argv=['atest', 'host', 'mod', '--protection',
-                               'Do not repair', 'host5' ,'--mlist', mfile.name,
-                               'host1', 'host6'],
-                         rpcs=[('modify_host', {'id': 'host6',
-                                                'protection': 'Do not repair'},
-                                True, None),
-                               ('modify_host', {'id': 'host5',
-                                                'protection': 'Do not repair'},
-                                True, None),
-                               ('modify_host', {'id': 'host4',
-                                                'protection': 'Do not repair'},
-                                True, None),
-                               ('modify_host', {'id': 'host3',
-                                                'protection': 'Do not repair'},
-                                True, None),
-                               ('modify_host', {'id': 'host2',
-                                                'protection': 'Do not repair'},
-                                True, None),
-                               ('modify_host', {'id': 'host1',
-                                                'protection': 'Do not repair'},
-                                True, None),
-                               ('modify_host', {'id': 'host0',
-                                                'protection': 'Do not repair'},
-                                True, None)],
-                         out_words_ok=['Do not repair', 'host0', 'host1',
-                                       'host2', 'host3', 'host4', 'host5',
-                                       'host6'])
+            self.run_cmd(argv=self._command_multiple + ['--unlock'], rpcs=rpcs,
+                         out_words_ok=out)
         finally:
             mfile.clean()
 
-    def test_execute_attribute_host(self):
-        self.run_cmd(argv=['atest', 'host', 'mod', 'host0', '--attribute',
-                           'foo=bar'],
-                     rpcs=[('modify_host', {'id': 'host0'}, True, None),
-                           ('set_host_attribute', {'hostname': 'host0',
-                                                   'attribute': 'foo',
-                                                   'value': 'bar'},
-                            True, None)],
-                     out_words_ok=[])
+
+    def test_single_attributes(self):
+        """Test applying one attribute to one host."""
+        attrs = {'foo': 'bar'}
+        s_attrs = ','.join(['='.join((k,v)) for k,v in attrs.items()])
+        rpcs, out = self._gen_expectations(attributes=attrs)
+        self.run_cmd(self._command_single + ['--attributes', s_attrs],
+                     rpcs=rpcs, out_words_ok=out)
 
 
-class host_create_unittest(cli_mock.cli_unittest):
-    _out = ['Added', 'host', 'localhost']
-    _command = ['atest', 'host', 'create', 'localhost']
+    def test_multiple_attributes_multiple_hosts(self):
+        """Test applying multiple attributes to multiple hosts."""
+        attrs = {'foo': 'bar', 'baz': 'zip'}
+        s_attrs = ','.join(['='.join((k,v)) for k,v in attrs.items()])
+        rpcs, out = self._gen_expectations(hosts=self._hosts, attributes=attrs)
+        self.run_cmd(self._command_multiple + ['--attributes', s_attrs],
+                     rpcs=rpcs, out_words_ok=out)
+
+
+    def test_platform(self):
+        """Test applying platform label."""
+        rpcs, out = self._gen_expectations(platform='some_platform')
+        self.run_cmd(argv=self._command_single + ['--platform',
+                                                  'some_platform'],
+                     rpcs=rpcs, out_words_ok=out)
+
+
+    def test_labels(self):
+        """Test applying labels."""
+        labels = ['label0', 'label1']
+        rpcs, out = self._gen_expectations(labels=labels)
+        self.run_cmd(argv=self._command_single + ['--labels', ','.join(labels)],
+                     rpcs=rpcs, out_words_ok=out)
+
+
+    def test_labels_from_file(self):
+        """Test applying labels from file."""
+        labels = ['label0', 'label1']
+        rpcs, out = self._gen_expectations(labels=labels)
+        labelsf = cli_mock.create_file(','.join(labels))
+        try:
+            self.run_cmd(argv=self._command_single + ['--blist', labelsf.name],
+                         rpcs=rpcs, out_words_ok=out)
+        finally:
+            labelsf.clean()
+
+
+    def test_acls(self):
+        """Test applying acls."""
+        acls = ['acl0', 'acl1']
+        rpcs, out = self._gen_expectations(acls=acls)
+        self.run_cmd(argv=self._command_single + ['--acls', ','.join(acls)],
+                     rpcs=rpcs, out_words_ok=out)
+
+
+    def test_acls_from_file(self):
+        """Test applying acls from file."""
+        acls = ['acl0', 'acl1']
+        rpcs, out = self._gen_expectations(acls=acls)
+        aclsf = cli_mock.create_file(','.join(acls))
+        try:
+            self.run_cmd(argv=self._command_single + ['-A', aclsf.name],
+                         rpcs=rpcs, out_words_ok=out)
+        finally:
+            aclsf.clean()
+
+
+    def test_protection(self):
+        """Test applying host protection."""
+        protection = 'Do not repair'
+        rpcs, out = self._gen_expectations(protection=protection)
+        self.run_cmd(argv=self._command_single + ['--protection', protection],
+                     rpcs=rpcs,out_words_ok=out)
+
+
+    def test_protection_invalid(self):
+        """Test invalid protection causes failure."""
+        protection = 'Invalid protection'
+        rpcs, out = self._gen_expectations(hosts=[])
+        self.run_cmd(argv=self._command_single + ['--protection', protection],
+                     exit_code=2, err_words_ok=['invalid', 'choice'] +
+                     protection.split())
+
+
+    def test_complex(self):
+        """Test applying multiple modifications / creating a complex host."""
+        lock_reason = 'Because I said so.'
+        platform = 'some_platform'
+        labels = ['label0', 'label1']
+        acls = ['acl0', 'acl1']
+        protection = 'Do not verify'
+        labelsf = cli_mock.create_file(labels[1])
+        aclsf = cli_mock.create_file(acls[1])
+        cmd_args = ['-l', '-r', lock_reason, '-t', platform, '-b', labels[0],
+                    '-B', labelsf.name, '-a', acls[0], '-A', aclsf.name, '-p',
+                    protection]
+        rpcs, out = self._gen_expectations(locked=True, lock_reason=lock_reason,
+                                           acls=acls, labels=labels,
+                                           platform=platform,
+                                           protection=protection)
+
+        try:
+            self.run_cmd(argv=self._command_single + cmd_args, rpcs=rpcs,
+                         out_words_ok=out)
+        finally:
+            labelsf.clean()
+            aclsf.clean()
+
+
+class host_mod_unittest(host_mod_create_tests, cli_mock.cli_unittest):
+    """Tests specific to the mod action and expectation generator for shared
+    tests.
+    """
+    _hosts = ['localhost', '127.0.0.1']
+    _host_ids = [1, 2]
+    _command_base = ['atest', 'host', 'mod']
+    _command_single = _command_base + [_hosts[0]]
+    _command_multiple = _command_base + _hosts
+
+    def _gen_expectations(self, hosts=['localhost'], locked=None,
+                          lock_reason='', force_lock=False, protection=None,
+                          acls=[], labels=[], platform=None, attributes={}):
+        rpcs = []
+        out = set()
+        hosts = hosts[:]
+        hosts.reverse()
+
+        # Genarate result for get_hosts command to include all known hosts
+        host_dicts = []
+        for h, h_id in zip(self._hosts, self._host_ids):
+            host_dicts.append({'hostname': h, 'id': h_id})
+        rpcs.append(('get_hosts', {'hostname__in': hosts}, True, host_dicts))
+
+        # Expect actions only for known hosts
+        host_ids = []
+        for host in hosts:
+            if host not in self._hosts:
+                continue
+            host_id = self._host_ids[self._hosts.index(host)]
+            host_ids.append(host_id)
+            modify_args = {'id': host}
+
+            if locked is not None:
+                out.add('Locked' if locked else 'Unlocked')
+                modify_args['locked'] = locked
+                modify_args['lock_reason'] = lock_reason
+            if force_lock:
+                modify_args['force_modify_locking'] = True
+            if protection:
+                modify_args['protection'] = protection
+
+            if len(modify_args.keys()) > 1:
+                out.add(host)
+                rpcs.append(('modify_host', modify_args, True, None))
+
+            if labels:
+                rpcs += self._gen_labels_rpcs(labels, host_id=host_id)
+                rpcs.append(('host_add_labels', {'id': host, 'labels': labels},
+                             True, None))
+
+            if platform:
+                rpcs += self._gen_labels_rpcs([platform], platform=True,
+                                              host_id=host_id)
+                rpcs.append(('host_add_labels', {'id': host,
+                                                 'labels': [platform]},
+                             True, None))
+
+            rpcs += self._gen_attributes_rpcs(host, attributes)
+
+        if acls:
+            rpcs += self._gen_acls_rpcs(hosts, acls, host_ids=host_ids)
+
+        return rpcs, list(out)
+
+
+    def test_mod_force_lock_one_host(self):
+        """Test mod with forced locking."""
+        lock_reason = 'Because'
+        rpcs, out = self._gen_expectations(locked=True, force_lock=True,
+                                           lock_reason=lock_reason)
+        self.run_cmd(argv=self._command_single + [
+                            '--lock', '--force_modify_locking', '--lock_reason',
+                            lock_reason],
+                     rpcs=rpcs, out_words_ok=out)
+
+    def test_mod_force_unlock_one_host(self):
+        """Test mod forced unlocking."""
+        rpcs, out = self._gen_expectations(locked=False, force_lock=True)
+        self.run_cmd(argv=self._command_single + ['--unlock',
+                                                   '--force_modify_locking'],
+                      rpcs=rpcs, out_words_ok=out)
+
+    def test_mod_fail_unknown_host(self):
+        """Test mod fails with unknown host."""
+        rpcs, out = self._gen_expectations(hosts=['nope'], locked=True)
+        self.run_cmd(argv=self._command_base + ['nope', '--lock'],
+                     rpcs=rpcs, err_words_ok=['Cannot', 'modify', 'nope'])
+
+
+class host_create_unittest(host_mod_create_tests, cli_mock.cli_unittest):
+    """Test specific to create action and expectation generator for shared
+    tests.
+    """
+    _hosts = ['localhost', '127.0.0.1']
+    _command_base = ['atest', 'host', 'create']
+    _command_single = _command_base + [_hosts[0]]
+    _command_multiple = _command_base + _hosts
 
     def _mock_host(self, platform=None, labels=[]):
         mock_host = self.god.create_mock_class(hosts.Host, 'Host')
         hosts.create_host = self.god.create_mock_function('create_host')
-        hosts.create_host.expect_any_call().and_return(mock_host)
-        mock_host.get_platform.expect_call().and_return(platform)
-        mock_host.get_labels.expect_call().and_return(labels)
         return mock_host
 
 
-    def _mock_testbed(self, platform=None, labels=[]):
-        mock_tb = self.god.create_mock_class(hosts.TestBed, 'TestBed')
-        hosts.create_testbed = self.god.create_mock_function('create_testbed')
-        hosts.create_testbed.expect_any_call().and_return(mock_tb)
-        mock_tb.get_platform.expect_call().and_return(platform)
-        mock_tb.get_labels.expect_call().and_return(labels)
-        return mock_tb
+    def _mock_create_host_call(self, mock_host, platform=None, labels=[]):
+        hosts.create_host.expect_any_call().and_return(mock_host)
+        mock_host.get_platform.expect_call().and_return(platform)
+        mock_host.get_labels.expect_call().and_return(labels)
 
 
-    def _gen_rpcs_for_label(self, label, platform=False):
-        rpcs = [
-            ('get_labels', {'name': label}, True, []),
-            ('add_label', {'name': label, 'platform': platform}, True, None)
-        ]
-        return rpcs
-
-
-    def _gen_expected_rpcs(self, hosts=None, locked=False,
-                           lock_reason=None, platform=None, labels=None,
-                           acls=None, protection=None, serials=None):
+    def _gen_expectations(self, hosts=['localhost'], locked=False,
+                          lock_reason=None, platform=None,
+                          discovered_platform=None, labels=[],
+                          discovered_labels=[], acls=[], protection=None,
+                          attributes={}):
         """Build a list of expected RPC calls based on values to host command.
 
         @param hosts: list of hostname being created (default ['localhost'])
         @param locked: end state of host (bool)
         @param lock_reason: reason for host to be locked
         @param platform: platform label
+        @param discovered_platform: platform discovered automatically by host
         @param labels: list of host labels (excluding platform)
+        @param discovered_labels: list of labels discovered automatically
         @param acls: list of host acls
         @param protection: host protection level
 
         @return: list of expect rpc calls (each call is (op, args, success,
             result))
         """
-        rpcs = []
-        hosts = hosts[:] if hosts else ['localhost']
+        hosts = hosts[:]
         hosts.reverse() # No idea why
         lock_reason = lock_reason or 'Forced lock on device creation'
         acls = acls or []
-        labels = labels or []
 
-        if platform:
-            rpcs += self._gen_rpcs_for_label(platform, platform=True)
-        for label in labels:
-            rpcs += self._gen_rpcs_for_label(label) * len(hosts)
+        rpcs = []
+        out = ['Added', 'host'] + hosts
 
-        for acl in acls:
-            rpcs.append(('get_acl_groups', {'name': acl}, True, []))
-            rpcs.append(('add_acl_group', {'name': acl}, True, None))
+        # Expect calls to create_host, host.labels and host.platform
+        mock_host = self._mock_host()
+        for host in hosts:
+            self._mock_create_host_call(mock_host, discovered_platform,
+                                        discovered_labels)
+
 
         for host in hosts:
             add_args = {
@@ -1443,41 +1632,20 @@
                 add_args['protection'] = protection
             rpcs.append(('add_host', add_args, True, None))
 
-            if labels or platform:
-                rpcs.append((
-                    'host_add_labels',
-                    {
-                        'id': host,
-                        'labels': labels + [platform] if platform else labels,
-                    },
-                    True,
-                    None
-                ))
+            rpcs += self._gen_labels_rpcs(labels)
+            if labels:
+                rpcs.append(('host_add_labels', {'id': host, 'labels': labels},
+                             True, None))
 
-        if serials:
-            for host in hosts:
-                rpcs.append((
-                    'set_host_attribute',
-                    {
-                        'hostname': host,
-                        'attribute': 'serials',
-                        'value': ','.join(serials),
-                    },
-                    True,
-                    None
-                ))
+            if platform:
+                rpcs += self._gen_labels_rpcs([platform], platform=True)
+                rpcs.append(('host_add_labels', {'id': host,
+                                                 'labels': [platform]},
+                             True, None))
 
-        for acl in acls:
-            for host in hosts:
-                rpcs.append((
-                    'acl_group_add_hosts',
-                    {
-                        'hosts': [host],
-                        'id': acl,
-                    },
-                    True,
-                    None,
-                ))
+            rpcs += self._gen_attributes_rpcs(host, attributes)
+
+        rpcs += self._gen_acls_rpcs(hosts, acls)
 
         if not locked:
             for host in hosts:
@@ -1491,165 +1659,46 @@
                     True,
                     None,
                 ))
-        return rpcs
+        return rpcs, out
+
+    def test_create_no_args(self):
+        """Test simple creation with to arguments."""
+        rpcs, out = self._gen_expectations()
+        self.run_cmd(argv=self._command_single, rpcs=rpcs, out_words_ok=out)
 
 
-    def test_create_simple(self):
-        self._mock_host()
-        rpcs = self._gen_expected_rpcs()
-        self.run_cmd(argv=self._command, rpcs=rpcs, out_words_ok=self._out)
-
-
-    def test_create_locked(self):
-        self._mock_host()
-        lock_reason = 'Because I said so.'
-        rpcs = self._gen_expected_rpcs(locked=True,
-                                                   lock_reason=lock_reason)
-        self.run_cmd(argv=self._command + ['-l', '-r', lock_reason],
-                     rpcs=rpcs, out_words_ok=self._out)
-
-
-    def test_create_discovered_platform(self):
-        self._mock_host(platform='some_platform')
-        rpcs = self._gen_expected_rpcs(platform='some_platform')
-        self.run_cmd(argv=self._command, rpcs=rpcs, out_words_ok=self._out)
-
-
-    def test_create_specified_platform(self):
-        self._mock_host()
-        rpcs = self._gen_expected_rpcs(platform='some_platform')
-        self.run_cmd(argv=self._command + ['-t', 'some_platform'], rpcs=rpcs,
-                     out_words_ok=self._out)
+    def test_create_with_discovered_platform(self):
+        """Test discovered platform is used when platform isn't specified."""
+        rpcs, out = self._gen_expectations(platform='some_platform',
+                                           discovered_platform='some_platform')
+        self.run_cmd(argv=self._command_single, rpcs=rpcs, out_words_ok=out)
 
 
     def test_create_specified_platform_overrides_discovered_platform(self):
-        self._mock_host(platform='wrong_platform')
-        rpcs = self._gen_expected_rpcs(platform='some_platform')
-        self.run_cmd(argv=self._command + ['-t', 'some_platform'], rpcs=rpcs,
-                     out_words_ok=self._out)
+        """Test that the specified platform overrides the discovered platform.
+        """
+        rpcs, out = self._gen_expectations(platform='some_platform',
+                                           discovered_platform='wrong_platform')
+        self.run_cmd(argv=self._command_single + ['--platform',
+                                                  'some_platform'],
+                     rpcs=rpcs, out_words_ok=out)
 
 
     def test_create_discovered_labels(self):
+        """Test applying automatically discovered labels."""
         labels = ['label0', 'label1']
-        self._mock_host(labels=labels)
-        rpcs = self._gen_expected_rpcs(labels=labels)
-        self.run_cmd(argv=self._command, rpcs=rpcs, out_words_ok=self._out)
+        rpcs, out = self._gen_expectations(labels=labels,
+                                           discovered_labels=labels)
+        self.run_cmd(argv=self._command_single, rpcs=rpcs, out_words_ok=out)
 
 
-    def test_create_specified_labels(self):
-        labels = ['label0', 'label1']
-        self._mock_host()
-        rpcs = self._gen_expected_rpcs(labels=labels)
-        self.run_cmd(argv=self._command + ['-b', ','.join(labels)], rpcs=rpcs,
-                     out_words_ok=self._out)
-
-
-    def test_create_specified_labels_from_file(self):
-        labels = ['label0', 'label1']
-        self._mock_host()
-        rpcs = self._gen_expected_rpcs(labels=labels)
-        labelsf = cli_mock.create_file(','.join(labels))
-        try:
-            self.run_cmd(argv=self._command + ['-B', labelsf.name], rpcs=rpcs,
-                         out_words_ok=self._out)
-        finally:
-            labelsf.clean()
-
     def test_create_specified_discovered_labels_combine(self):
+        """Test applying both discovered and specified labels."""
         labels = ['label0', 'label1']
-        self._mock_host(labels=labels[0:1])
-        rpcs = self._gen_expected_rpcs(labels=labels)
-        self.run_cmd(argv=self._command + ['-b', labels[1]], rpcs=rpcs,
-                     out_words_ok=self._out)
-
-
-    def test_create_acls(self):
-        acls = ['acl0', 'acl1']
-        self._mock_host()
-        rpcs = self._gen_expected_rpcs(acls=acls)
-        self.run_cmd(argv=self._command + ['-a', ','.join(acls)], rpcs=rpcs,
-                     out_words_ok=self._out)
-
-
-    def test_create_acls_from_file(self):
-        acls = ['acl0', 'acl1']
-        self._mock_host()
-        rpcs = self._gen_expected_rpcs(acls=acls)
-        aclsf = cli_mock.create_file(','.join(acls))
-        try:
-            self.run_cmd(argv=self._command + ['-A', aclsf.name], rpcs=rpcs,
-                         out_words_ok=self._out)
-        finally:
-            aclsf.clean()
-
-
-    def test_create_protection(self):
-        protection = 'Do not repair'
-        self._mock_host()
-        rpcs = self._gen_expected_rpcs(protection=protection)
-        self.run_cmd(argv=self._command + ['-p', protection], rpcs=rpcs,
-                     out_words_ok=self._out)
-
-
-    def test_create_protection_invalid(self):
-        protection = 'Invalid protection'
-        rpcs = self._gen_expected_rpcs()
-        self.run_cmd(argv=self._command + ['-p', protection], exit_code=2,
-                     err_words_ok=['invalid', 'choice'] + protection.split())
-
-
-    def test_create_one_serial(self):
-        serial = 'device0'
-        self._mock_host()
-        rpcs = self._gen_expected_rpcs(serials=[serial])
-        self.run_cmd(argv=self._command + ['-s', serial], rpcs=rpcs,
-                     out_words_ok=self._out)
-
-
-    def test_create_multiple_serials(self):
-        serials = ['device0', 'device1']
-        self._mock_testbed()
-        rpcs = self._gen_expected_rpcs(serials=serials)
-        self.run_cmd(argv=self._command + ['-s', ','.join(serials)], rpcs=rpcs,
-                     out_words_ok=self._out)
-
-
-    def test_create_multiple_simple_hosts(self):
-        mock_host = self._mock_host()
-        hosts.create_host.expect_any_call().and_return(mock_host)
-        mock_host.get_platform.expect_call()
-        mock_host.get_labels.expect_call().and_return([])
-
-        hostnames = ['localhost', '127.0.0.1']
-        rpcs = self._gen_expected_rpcs(hosts=hostnames)
-
-        self.run_cmd(argv=['atest', 'host', 'create'] + hostnames,
-                     rpcs=rpcs[0:4],
-                     out_words_ok=['Added', 'hosts'] + hostnames)
-
-
-    def test_create_complex(self):
-        lock_reason = 'Because I said so.'
-        platform = 'some_platform'
-        labels = ['label0', 'label1', 'label2']
-        acls = ['acl0', 'acl1']
-        protection = 'Do not verify'
-        labelsf = cli_mock.create_file(labels[2])
-        aclsf = cli_mock.create_file(acls[1])
-        cmd_args = ['-l', '-r', lock_reason, '-t', platform, '-b', labels[1],
-                    '-B', labelsf.name, '-a', acls[0], '-A', aclsf.name, '-p',
-                    protection]
-        self._mock_host(labels=labels[0:1])
-        rpcs = self._gen_expected_rpcs(locked=True, lock_reason=lock_reason,
-                                       acls=acls, labels=labels,
-                                       platform=platform, protection=protection)
-
-        try:
-            self.run_cmd(argv=self._command + cmd_args, rpcs=rpcs,
-                         out_words_ok=self._out)
-        finally:
-            labelsf.clean()
-            aclsf.clean()
+        rpcs, out = self._gen_expectations(labels=labels,
+                                           discovered_labels=[labels[0]])
+        self.run_cmd(argv=self._command_single + ['--labels', labels[1]],
+                     rpcs=rpcs, out_words_ok=out)
 
 
 if __name__ == '__main__':
diff --git a/cli/site_host.py b/cli/site_host.py
deleted file mode 100644
index ca8a798..0000000
--- a/cli/site_host.py
+++ /dev/null
@@ -1,184 +0,0 @@
-# Copyright (c) 2011 The Chromium OS Authors. All rights reserved.
-# Use of this source code is governed by a BSD-style license that can be
-# found in the LICENSE file.
-
-import common
-import inspect, new, socket, sys
-
-from autotest_lib.client.bin import utils
-from autotest_lib.cli import host, rpc
-from autotest_lib.server import hosts
-from autotest_lib.client.common_lib import error, host_protections
-
-
-# In order for hosts to work correctly, some of its variables must be setup.
-hosts.factory.ssh_user = 'root'
-hosts.factory.ssh_pass = ''
-hosts.factory.ssh_port = 22
-hosts.factory.ssh_verbosity_flag = ''
-hosts.factory.ssh_options = ''
-
-
-class site_host(host.host):
-    """Required by atest's site logic."""
-    pass
-
-
-class site_host_create(site_host, host.host_create):
-    """
-    site_host_create subclasses host_create in host.py.
-    """
-
-    @classmethod
-    def construct_without_parse(
-            cls, web_server, hosts, platform=None,
-            locked=False, lock_reason='', labels=[], acls=[],
-            protection=host_protections.Protection.NO_PROTECTION):
-        """Construct an site_host_create object and fill in data from args.
-
-        Do not need to call parse after the construction.
-
-        Return an object of site_host_create ready to execute.
-
-        @param web_server: A string specifies the autotest webserver url.
-            It is needed to setup comm to make rpc.
-        @param hosts: A list of hostnames as strings.
-        @param platform: A string or None.
-        @param locked: A boolean.
-        @param lock_reason: A string.
-        @param labels: A list of labels as strings.
-        @param acls: A list of acls as strings.
-        @param protection: An enum defined in host_protections.
-        """
-        obj = cls()
-        obj.web_server = web_server
-        try:
-            # Setup stuff needed for afe comm.
-            obj.afe = rpc.afe_comm(web_server)
-        except rpc.AuthError, s:
-            obj.failure(str(s), fatal=True)
-        obj.hosts = hosts
-        obj.platform = platform
-        obj.locked = locked
-        if locked and lock_reason.strip():
-            obj.data['lock_reason'] = lock_reason.strip()
-        obj.labels = labels
-        obj.acls = acls
-        if protection:
-            obj.data['protection'] = protection
-        # TODO(kevcheng): Update the admin page to take in serials?
-        obj.serials = None
-        return obj
-
-
-    def _execute_add_one_host(self, host):
-        # Always add the hosts as locked to avoid the host
-        # being picked up by the scheduler before it's ACL'ed.
-        self.data['locked'] = True
-        if not self.locked:
-            self.data['lock_reason'] = 'Forced lock on device creation'
-        self.execute_rpc('add_host', hostname=host,
-                         status="Ready", **self.data)
-        # If there are labels avaliable for host, use them.
-        host_info = self.host_info_map[host]
-        labels = set(self.labels)
-        if host_info.labels:
-            labels.update(host_info.labels)
-        # Now add the platform label.
-        # If a platform was not provided and we were able to retrieve it
-        # from the host, use the retrieved platform.
-        platform = self.platform if self.platform else host_info.platform
-        if platform:
-            labels.add(platform)
-
-        if len(labels):
-            self.execute_rpc('host_add_labels', id=host, labels=list(labels))
-
-        if self.serials:
-            self.execute_rpc('set_host_attribute', hostname=host,
-                             attribute='serials', value=','.join(self.serials))
-
-    def execute(self):
-        # Check to see if the platform or any other labels can be grabbed from
-        # the hosts.
-        self.host_info_map = {}
-        for host in self.hosts:
-            try:
-                if utils.ping(host, tries=1, deadline=1) == 0:
-                    if self.serials and len(self.serials) > 1:
-                        host_dut = hosts.create_testbed(
-                                host, adb_serials=self.serials)
-                    else:
-                        adb_serial = None
-                        if self.serials:
-                            adb_serial = self.serials[0]
-                        host_dut = hosts.create_host(host,
-                                                     adb_serial=adb_serial)
-                    host_info = host_information(host,
-                                                 host_dut.get_platform(),
-                                                 host_dut.get_labels())
-                else:
-                    # Can't ping the host, use default information.
-                    host_info = host_information(host, None, [])
-            except (socket.gaierror, error.AutoservRunError,
-                    error.AutoservSSHTimeout):
-                # We may be adding a host that does not exist yet or we can't
-                # reach due to hostname/address issues or if the host is down.
-                host_info = host_information(host, None, [])
-            self.host_info_map[host] = host_info
-        # We need to check if these labels & ACLs exist,
-        # and create them if not.
-        if self.platform:
-            self.check_and_create_items('get_labels', 'add_label',
-                                        [self.platform],
-                                        platform=True)
-        else:
-            # No platform was provided so check and create the platform label
-            # for each host.
-            platforms = []
-            for host_info in self.host_info_map.values():
-                if host_info.platform and host_info.platform not in platforms:
-                    platforms.append(host_info.platform)
-            if platforms:
-                self.check_and_create_items('get_labels', 'add_label',
-                                            platforms,
-                                            platform=True)
-        labels_to_check_and_create = self.labels[:]
-        for host_info in self.host_info_map.values():
-            labels_to_check_and_create = (host_info.labels +
-                                          labels_to_check_and_create)
-        if labels_to_check_and_create:
-            self.check_and_create_items('get_labels', 'add_label',
-                                        labels_to_check_and_create,
-                                        platform=False)
-
-        if self.acls:
-            self.check_and_create_items('get_acl_groups',
-                                        'add_acl_group',
-                                        self.acls)
-
-        return self._execute_add_hosts()
-
-
-class host_information(object):
-    """Store host information so we don't have to keep looking it up."""
-
-
-    def __init__(self, hostname, platform, labels):
-        self.hostname = hostname
-        self.platform = platform
-        self.labels = labels
-
-
-# Any classes we don't override in host should be copied automatically
-for cls in [getattr(host, n) for n in dir(host) if not n.startswith("_")]:
-    if not inspect.isclass(cls):
-        continue
-    cls_name = cls.__name__
-    site_cls_name = 'site_' + cls_name
-    if hasattr(sys.modules[__name__], site_cls_name):
-        continue
-    bases = (site_host, cls)
-    members = {'__doc__': cls.__doc__}
-    site_cls = new.classobj(site_cls_name, bases, members)
-    setattr(sys.modules[__name__], site_cls_name, site_cls)
diff --git a/frontend/afe/admin.py b/frontend/afe/admin.py
index 5b7143e..72ab046 100644
--- a/frontend/afe/admin.py
+++ b/frontend/afe/admin.py
@@ -7,7 +7,7 @@
 from django.utils.encoding import smart_str
 from django.utils.safestring import mark_safe
 
-from autotest_lib.cli import rpc, site_host
+from autotest_lib.cli import rpc, host
 from autotest_lib.frontend import settings
 from autotest_lib.frontend.afe import model_logic, models
 
@@ -227,7 +227,7 @@
         acls = []
 
         # Pipe to cli to perform autodetection and create host.
-        host_create_obj = site_host.site_host_create.construct_without_parse(
+        host_create_obj = shost.host_create.construct_without_parse(
                 web_server, hosts, platform,
                 locked, lock_reason, labels, acls,
                 protection)