[autotest] Use HostInfo to update version labels

This CL slightly changes the semantics of clear_version_labels: All
versions labels are cleared now instead of just the type of version
recognized by the corresponding host class. This was always the intended
behaviour.

This CL also changes the semantics from afe_utils.add_version_label in
that if the version prefix already exists, it is updated in place. This
is the right thing to do because adding a version label to the end of
the list when an earlier label exists has no effect (the first one takes
precedence).

BUG=chromium:678430
TEST=(1) (new) unittests
     (2) in-lab hwtests via moblab

Change-Id: I3939e04953a1dd2ac0cf09f61462128410c0db06
Reviewed-on: https://chromium-review.googlesource.com/517772
Commit-Ready: Prathmesh Prabhu <pprabhu@chromium.org>
Tested-by: Prathmesh Prabhu <pprabhu@chromium.org>
Reviewed-by: Prathmesh Prabhu <pprabhu@chromium.org>
diff --git a/server/afe_utils.py b/server/afe_utils.py
index 1466cef..c40e31c 100644
--- a/server/afe_utils.py
+++ b/server/afe_utils.py
@@ -197,7 +197,9 @@
     @param *args: Args list to pass to machine_install.
     @param **dargs: dargs dict to pass to machine_install.
     """
-    clear_version_labels(host)
+    info = host.host_info_store.get()
+
+    info.clear_version_labels()
     clear_host_attributes_before_provision(host)
     # If ENABLE_DEVSERVER_TRIGGER_AUTO_UPDATE is enabled and the host is a
     # CrosHost, devserver will be used to trigger auto-update.
@@ -208,4 +210,6 @@
         image_name, host_attributes = host.machine_install(*args, **dargs)
     for attribute, value in host_attributes.items():
         update_host_attribute(host, attribute, value)
-    add_version_label(host, image_name)
+    info.set_version_label(host.VERSION_PREFIX, image_name)
+
+    host.host_info_store.commit(info)
diff --git a/server/hosts/afe_store.py b/server/hosts/afe_store.py
index 71b07c2..a41cf4a 100644
--- a/server/hosts/afe_store.py
+++ b/server/hosts/afe_store.py
@@ -61,8 +61,10 @@
         # based on that. If another user tries to commit it's changes in
         # parallel, we'll end up with corrupted labels / attributes.
         old_info = self._refresh_impl()
-        self._remove_labels_on_afe(set(old_info.labels) - set(new_info.labels))
-        self._add_labels_on_afe(set(new_info.labels) - set(old_info.labels))
+        self._remove_labels_on_afe(
+                list(set(old_info.labels) - set(new_info.labels)))
+        self._add_labels_on_afe(
+                list(set(new_info.labels) - set(old_info.labels)))
 
         # TODO(pprabhu) Also commit attributes when we first replace a direct
         # AFE call for attribute update.
diff --git a/server/hosts/afe_store_unittest.py b/server/hosts/afe_store_unittest.py
index ed0f638..c24cae3 100644
--- a/server/hosts/afe_store_unittest.py
+++ b/server/hosts/afe_store_unittest.py
@@ -68,9 +68,9 @@
         self.assertEqual(self.mock_afe.run.call_count, 2)
         expected_run_calls = [
                 mock.call('host_remove_labels', id='some-host',
-                          labels={'label1'}),
+                          labels=['label1']),
                 mock.call('host_add_labels', id='some-host',
-                          labels={'label2'}),
+                          labels=['label2']),
         ]
         self.mock_afe.run.assert_has_calls(expected_run_calls,
                                            any_order=True)
diff --git a/server/hosts/host_info.py b/server/hosts/host_info.py
index 3498d1d..bb4db09 100644
--- a/server/hosts/host_info.py
+++ b/server/hosts/host_info.py
@@ -37,6 +37,12 @@
     _OS_PREFIX = 'os'
     _POOL_PREFIX = 'pool'
 
+    _VERSION_LABELS = (
+            provision.CROS_VERSION_PREFIX,
+            provision.ANDROID_BUILD_VERSION_PREFIX,
+            provision.TESTBED_BUILD_VERSION_PREFIX,
+    )
+
     def __init__(self, labels=None, attributes=None):
         """
         @param labels: (optional list) labels to set on the HostInfo.
@@ -56,9 +62,7 @@
         @returns The first build label for this host (if there are multiple).
                 None if no build label is found.
         """
-        for label_prefix in [provision.CROS_VERSION_PREFIX,
-                            provision.ANDROID_BUILD_VERSION_PREFIX,
-                            provision.TESTBED_BUILD_VERSION_PREFIX]:
+        for label_prefix in self._VERSION_LABELS:
             build_labels = self._get_stripped_labels_with_prefix(label_prefix)
             if build_labels:
                 return build_labels[0]
@@ -104,6 +108,33 @@
         return values[0] if values else ''
 
 
+    def clear_version_labels(self):
+        """Clear all version labels for the host."""
+        self.labels = [
+                label for label in self.labels if
+                not any(label.startswith(prefix + ':')
+                        for prefix in self._VERSION_LABELS)]
+
+
+    def set_version_label(self, version_prefix, version):
+        """Sets the version label for the host.
+
+        If a label with version_prefix exists, this updates the value for that
+        label, else appends a new label to the end of the label list.
+
+        @param version_prefix: The prefix to use (without the infix ':').
+        @param version: The version label value to set.
+        """
+        full_prefix = _to_label_prefix(version_prefix)
+        new_version_label = full_prefix + version
+        for index, label in enumerate(self.labels):
+            if label.startswith(full_prefix):
+                self.labels[index] = new_version_label
+                return
+        else:
+            self.labels.append(new_version_label)
+
+
     def _get_stripped_labels_with_prefix(self, prefix):
         """Search for labels with the prefix and remove the prefix.
 
@@ -340,3 +371,12 @@
                         deserialized_json['attributes'])
     except KeyError as e:
         raise DeserializationError('Malformed serialized host_info: %r' % e)
+
+
+def _to_label_prefix(prefix):
+    """Ensure that prefix has the expected format for label prefixes.
+
+    @param prefix: The (str) prefix to sanitize.
+    @returns: The sanitized (str) prefix.
+    """
+    return prefix if prefix.endswith(':') else prefix + ':'
diff --git a/server/hosts/host_info_unittest.py b/server/hosts/host_info_unittest.py
index da7387b..40ed128 100644
--- a/server/hosts/host_info_unittest.py
+++ b/server/hosts/host_info_unittest.py
@@ -151,6 +151,48 @@
                          "HostInfo [Labels: ['a'], Attributes: {'b': 2}]")
 
 
+    def test_clear_version_labels_no_labels(self):
+        """When no version labels exit, do nothing for clear_version_labels."""
+        original_labels = ['board:something', 'os:something_else',
+                           'pool:mypool', 'ab-version-corrupted:blah',
+                           'cros-version']
+        self.info.labels = list(original_labels)
+        self.info.clear_version_labels()
+        self.assertListEqual(self.info.labels, original_labels)
+
+
+    def test_clear_all_version_labels(self):
+        """Clear each recognized type of version label."""
+        original_labels = ['extra_label', 'cros-version:cr1', 'ab-version:ab1',
+                           'testbed-version:tb1']
+        self.info.labels = list(original_labels)
+        self.info.clear_version_labels()
+        self.assertListEqual(self.info.labels, ['extra_label'])
+
+    def test_clear_all_version_label_prefixes(self):
+        """Clear each recognized type of version label with empty value."""
+        original_labels = ['extra_label', 'cros-version:', 'ab-version:',
+                           'testbed-version:']
+        self.info.labels = list(original_labels)
+        self.info.clear_version_labels()
+        self.assertListEqual(self.info.labels, ['extra_label'])
+
+
+    def test_set_version_labels_updates_in_place(self):
+        """Update version label in place if prefix already exists."""
+        self.info.labels = ['extra', 'cros-version:X', 'ab-version:Y']
+        self.info.set_version_label('cros-version', 'Z')
+        self.assertListEqual(self.info.labels, ['extra', 'cros-version:Z',
+                                                'ab-version:Y'])
+
+    def test_set_version_labels_appends(self):
+        """Append a new version label if the prefix doesn't exist."""
+        self.info.labels = ['extra', 'ab-version:Y']
+        self.info.set_version_label('cros-version', 'Z')
+        self.assertListEqual(self.info.labels, ['extra', 'ab-version:Y',
+                                                'cros-version:Z'])
+
+
 class InMemoryHostInfoStoreTest(unittest.TestCase):
     """Basic tests for CachingHostInfoStore using InMemoryHostInfoStore."""