[autotest] Standardize label logic

BUG=None
TEST=Run dummy suite with provisioning locally and run unittest

Change-Id: Ie71009972ea950e1d8402bbbb12e0172c43baffa
Reviewed-on: https://chromium-review.googlesource.com/400618
Commit-Ready: Allen Li <ayatane@chromium.org>
Tested-by: Allen Li <ayatane@chromium.org>
Reviewed-by: Richard Barnette <jrbarnette@google.com>
diff --git a/server/cros/dynamic_suite/dynamic_suite.py b/server/cros/dynamic_suite/dynamic_suite.py
index 805a010..3f48e46 100644
--- a/server/cros/dynamic_suite/dynamic_suite.py
+++ b/server/cros/dynamic_suite/dynamic_suite.py
@@ -469,7 +469,7 @@
     # version_prefix+build should make it into each test as a DEPENDENCY.  The
     # easiest way to do this is to tack it onto the suite_dependencies.
     suite_spec.suite_dependencies.extend(
-            provision.join(version_prefix, build)
+            provision.NamespaceLabel(version_prefix, build)
             for version_prefix, build in suite_spec.builds.items())
 
     afe = frontend_wrappers.RetryingAFE(timeout_min=30, delay_sec=10,
diff --git a/server/cros/provision.py b/server/cros/provision.py
index 5e7608d..a6c3cd9 100644
--- a/server/cros/provision.py
+++ b/server/cros/provision.py
@@ -2,7 +2,8 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
-
+import collections
+import re
 import sys
 
 import common
@@ -24,38 +25,212 @@
 FLAKY_DEVSERVER_ATTEMPTS = 2
 
 
-### Helpers to convert value to label
-def cros_version_to_label(image):
+def label_from_str(label_string):
+    """Return a proper Label instance from a label string.
+
+    This function is for converting an existing label string into a Label
+    instance of the proper type.  For constructing a specific type of label,
+    don't use this.  Instead, instantiate the specific Label subclass.
+
+    @param label_string: Label string.
+    @returns: An instance of Label or a subclass.
     """
-    Returns the proper label name for a ChromeOS build of |image|.
+    if NamespaceLabel.SEP in label_string:
+        label = NamespaceLabel.from_str(label_string)
+        namespaces = _PresetNamespaceLabelMeta.namespaces
+        if label.namespace in namespaces:
+            return namespaces[label.namespace](label.value)
+        else:
+            return label
+    else:
+        return Label(label_string)
 
-    @param image: A string of the form 'lumpy-release/R28-3993.0.0'
-    @returns: A string that is the appropriate label name.
 
+class Label(str):
+    """A string that is explicitly typed as a label."""
+
+    def __repr__(self):
+        return '{cls}({label})'.format(
+            cls=type(self).__name__,
+            label=super(Label, self).__repr__())
+
+    @property
+    def action(self):
+        """Return the action represented by the label.
+
+        This is used for determine actions to perform based on labels, for
+        example for provisioning or repair.
+
+        @return: An Action instance.
+        """
+        return Action(self, '')
+
+
+Action = collections.namedtuple('Action', 'name,value')
+
+
+class NamespaceLabel(Label):
+    """Label with namespace and value separated by a colon."""
+
+    SEP = ':'
+
+    def __new__(cls, namespace, value):
+        return super(NamespaceLabel, cls).__new__(
+            cls, cls.SEP.join((namespace, value)))
+
+    @classmethod
+    def from_str(cls, label):
+        """Make NamespaceLabel instance from full string.
+
+        @param label: Label string.
+        @returns: NamespaceLabel instance.
+        """
+        namespace, value = label.split(cls.SEP, 1)
+        return cls(namespace, value)
+
+    @property
+    def namespace(self):
+        """The label's namespace (before colon).
+
+        @returns: string
+        """
+        return self.split(self.SEP, 1)[0]
+
+    @property
+    def value(self):
+        """The label's value (after colon).
+
+        @returns: string
+        """
+        return self.split(self.SEP, 1)[1]
+
+    @property
+    def action(self):
+        """Return the action represented by the label.
+
+        See docstring on overridden method.
+
+        @return: An Action instance.
+        """
+        return Action(self.namespace, self.value)
+
+
+class _PresetNamespaceLabelMeta(type):
+    """Metaclass for PresetNamespaceLabelMeta and subclasses.
+
+    This automatically tracks the NAMESPACE for concrete classes that define
+    it.  The namespaces attribute is a dict mapping namespace strings to the
+    corresponding NamespaceLabel subclass.
     """
-    return CROS_VERSION_PREFIX + ':' + image
+
+    namespaces = {}
+
+    def __init__(cls, name, bases, dict_):
+        if hasattr(cls, 'NAMESPACE'):
+            type(cls).namespaces[cls.NAMESPACE] = cls
 
 
-def fwro_version_to_label(image):
+class _PresetNamespaceLabel(NamespaceLabel):
+    """NamespaceLabel with preset namespace.
+
+    This class is abstract.  Concrete subclasses must set a NAMESPACE class
+    attribute.
     """
-    Returns the proper label name for a RO firmware build of |image|.
 
-    @param image: A string of the form 'lumpy-release/R28-3993.0.0'
-    @returns: A string that is the appropriate label name.
+    __metaclass__ = _PresetNamespaceLabelMeta
 
+    def __new__(cls, value):
+        return super(_PresetNamespaceLabel, cls).__new__(cls, cls.NAMESPACE, value)
+
+
+class CrosVersionLabel(_PresetNamespaceLabel):
+    """cros-version label."""
+    NAMESPACE = CROS_VERSION_PREFIX
+
+    @property
+    def value(self):
+        """The label's value (after colon).
+
+        @returns: string
+        """
+        return CrosVersion(super(CrosVersionLabel, self).value)
+
+
+class FWROVersionLabel(_PresetNamespaceLabel):
+    """Read-only firmware version label."""
+    NAMESPACE = FW_RO_VERSION_PREFIX
+
+
+class FWRWVersionLabel(_PresetNamespaceLabel):
+    """Read-write firmware version label."""
+    NAMESPACE = FW_RW_VERSION_PREFIX
+
+
+class CrosVersion(str):
+    """The name of a CrOS image version (e.g. lumpy-release/R27-3773.0.0).
+
+    Parts of the image name are exposed via properties.  In case the name is
+    not well-formed, these properties return INVALID_STR, which is not a valid value
+    for any part.
+
+    Class attributes:
+        INVALID_STR -- String returned if the version is not well-formed.
+
+    Properties:
+        group
+        milestone
+        version
+        rc
     """
-    return FW_RO_VERSION_PREFIX + ':' + image
 
+    INVALID_STR = 'N/A'
+    _NAME_PATTERN = re.compile(
+        r'^'
+        r'(?P<group>[a-z0-9-]+)'
+        r'/'
+        r'(?P<milestone>R[0-9]+)'
+        r'-'
+        r'(?P<version>[0-9.]+)'
+        r'(-(?P<rc>rc[0-9]+))?'
+        r'$'
+    )
 
-def fwrw_version_to_label(image):
-    """
-    Returns the proper label name for a RW firmware build of |image|.
+    def __repr__(self):
+        return '{cls}({name})'.format(
+            cls=type(self).__name__,
+            name=super(CrosVersion, self).__repr__())
 
-    @param image: A string of the form 'lumpy-release/R28-3993.0.0'
-    @returns: A string that is the appropriate label name.
+    def _get_group(self, group):
+        """Get regex match group, or fall back to N/A.
 
-    """
-    return FW_RW_VERSION_PREFIX + ':' + image
+        @param group: Group name string.
+        @returns String.
+        """
+        match = self._NAME_PATTERN.search(self)
+        if match is None:
+            return self.INVALID_STR
+        else:
+            return match.group(group)
+
+    @property
+    def group(self):
+        """Cros image group (e.g. lumpy-release)."""
+        return self._get_group('group')
+
+    @property
+    def milestone(self):
+        """Cros image milestone (e.g. R27)."""
+        return self._get_group('milestone')
+
+    @property
+    def version(self):
+        """Cros image milestone (e.g. 3773.0.0)."""
+        return self._get_group('version')
+
+    @property
+    def rc(self):
+        """Cros image rc (e.g. rc2)."""
+        return self._get_group('rc')
 
 
 class _SpecialTaskAction(object):
@@ -77,19 +252,17 @@
     # across available label prefixes.
     _priorities = []
 
-
     @classmethod
-    def acts_on(cls, label):
+    def acts_on(cls, label_string):
         """
         Returns True if the label is a label that we recognize as something we
         know how to act on, given our _actions.
 
-        @param label: The label as a string.
+        @param label_string: The label as a string.
         @returns: True if there exists a test to run for this label.
-
         """
-        return label.split(':')[0] in cls._actions
-
+        label = label_from_str(label_string)
+        return label.action.name in cls._actions
 
     @classmethod
     def test_for(cls, label):
@@ -99,11 +272,9 @@
         @param label: The label for which the action is being requested.
         @returns: The string name of the test that should be run.
         @raises KeyError: If the name was not recognized as one we care about.
-
         """
         return cls._actions[label]
 
-
     @classmethod
     def partition(cls, labels):
         """
@@ -130,30 +301,31 @@
 
         return capabilities, configurations
 
-
     @classmethod
-    def sort_configurations(cls, configurations):
+    def get_sorted_actions(cls, configurations):
         """
         Sort configurations based on the priority defined in cls._priorities.
 
         @param configurations: A list of actionable labels.
-
-        @return: A sorted list of tuple of (label_prefix, value), the tuples are
-                sorted based on the label_prefix's index in cls._priorities.
+        @return: A list of Action instances sorted by the action name in
+            cls._priorities.
         """
-        # Split a list of labels into a dict mapping name to value.  All labels
-        # must be provisionable labels, or else a ValueError
-        # For example, label 'cros-version:lumpy-release/R28-3993.0.0' is split
-        # to  {'cros-version': 'lumpy-release/R28-3993.0.0'}
-        split_configurations = dict()
-        for label in configurations:
-            name, _, value = label.partition(':')
-            split_configurations[name] = value
+        actions = (label_from_str(label_string).action
+                   for label_string in configurations)
+        return sorted(actions, key=cls._get_action_priority)
 
-        sort_key = (lambda config:
-                (cls._priorities.index(config[0])
-                 if (config[0] in cls._priorities) else sys.maxint))
-        return sorted(split_configurations.items(), key=sort_key)
+    @classmethod
+    def _get_action_priority(cls, action):
+        """
+        Return the priority of the action string.
+
+        @param action: An Action instance.
+        @return: An int.
+        """
+        if action.name in cls._priorities:
+            return cls._priorities.index(action.name)
+        else:
+            return sys.maxint
 
 
 class Verify(_SpecialTaskAction):
@@ -266,21 +438,6 @@
             label == SKIP_PROVISION)
 
 
-def join(provision_type, provision_value):
-    """
-    Combine the provision type and value into the label name.
-
-    @param provision_type: One of the constants that are the label prefixes.
-    @param provision_value: A string of the value for this provision type.
-    @returns: A string that is the label name for this (type, value) pair.
-
-    >>> join(CROS_VERSION_PREFIX, 'lumpy-release/R27-3773.0.0')
-    'cros-version:lumpy-release/R27-3773.0.0'
-
-    """
-    return '%s:%s' % (provision_type, provision_value)
-
-
 class SpecialTaskActionException(Exception):
     """
     Exception raised when a special task fails to successfully run a test that
@@ -312,8 +469,8 @@
                    "Can't %s label '%s'." % (task.name, label))
 
     # Sort the configuration labels based on `task._priorities`.
-    sorted_configurations = task.sort_configurations(configurations)
-    for name, value in sorted_configurations:
+    actions = task.get_sorted_actions(configurations)
+    for name, value in actions:
         action_item = task.test_for(name)
         success = action_item.execute(job=job, host=host, value=value)
         if not success:
diff --git a/server/cros/provision_unittest.py b/server/cros/provision_unittest.py
new file mode 100755
index 0000000..0d9c1bd
--- /dev/null
+++ b/server/cros/provision_unittest.py
@@ -0,0 +1,184 @@
+#!/usr/bin/python
+#
+# Copyright (c) 2016 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 unittest
+
+import common
+from autotest_lib.server.cros import provision
+
+
+class LabelFromStrTestCase(unittest.TestCase):
+    """label_from_str() test case."""
+
+    def test_fallback_label_unchanged(self):
+        """Test that Label doesn't change str value."""
+        label_str = 'dummy_label'
+        got = provision.label_from_str(label_str)
+        self.assertEqual(got, label_str)
+
+    def test_fallback_label_type(self):
+        """Test that label_from_str() falls back to Label."""
+        label_str = 'dummy_label'
+        got = provision.label_from_str(label_str)
+        self.assertIsInstance(got, provision.Label)
+
+    def test_fallback_namespace_unchanged(self):
+        """Test that NamespaceLabel doesn't change str value."""
+        label_str = 'dummy-namespace:value'
+        got = provision.label_from_str(label_str)
+        self.assertEqual(got, label_str)
+
+    def test_fallback_namespace_label_type(self):
+        """Test that label_from_str() falls back to NamespaceLabel."""
+        label_str = 'dummy-namespace:value'
+        got = provision.label_from_str(label_str)
+        self.assertIsInstance(got, provision.NamespaceLabel)
+
+    def test_cros_version_unchanged(self):
+        """Test that CrosVersionLabel doesn't change str value."""
+        label_str = 'cros-version:value'
+        got = provision.label_from_str(label_str)
+        self.assertEqual(got, label_str)
+
+    def test_cros_version_label_type(self):
+        """Test that label_from_str() detects cros-version."""
+        label_str = 'cros-version:value'
+        got = provision.label_from_str(label_str)
+        self.assertIsInstance(got, provision.CrosVersionLabel)
+
+    def test_fwrw_version_label_type(self):
+        """Test that label_from_str() detects fwrw-version."""
+        label_str = 'fwrw-version:value'
+        got = provision.label_from_str(label_str)
+        self.assertIsInstance(got, provision.FWRWVersionLabel)
+
+    def test_fwro_version_label_type(self):
+        """Test that label_from_str() detects fwro-version."""
+        label_str = 'fwro-version:value'
+        got = provision.label_from_str(label_str)
+        self.assertIsInstance(got, provision.FWROVersionLabel)
+
+
+class LabelTestCase(unittest.TestCase):
+    """Label test case."""
+
+    def test_label_repr(self):
+        """Test that Label repr works."""
+        label = provision.Label('dummy_label')
+        self.assertEqual(repr(label), "Label('dummy_label')")
+
+    def test_label_eq_str(self):
+        """Test that Label equals its string value."""
+        label = provision.Label('dummy_label')
+        self.assertEqual(label, 'dummy_label')
+
+    def test_get_action(self):
+        """Test Label action property."""
+        action = provision.Label('dummy_label').action
+        self.assertEqual(action, provision.Action('dummy_label', ''))
+
+
+class NamespaceLabelTestCase(unittest.TestCase):
+    """NamespaceLabel test case."""
+
+    def test_get_namespace(self):
+        """Test NamespaceLabel namespace property."""
+        label = provision.NamespaceLabel('ns', 'value')
+        self.assertEqual(label.namespace, 'ns')
+
+    def test_get_value(self):
+        """Test NamespaceLabel value property."""
+        label = provision.NamespaceLabel('ns', 'value')
+        self.assertEqual(label.value, 'value')
+
+    def test_from_str_identity(self):
+        """Test NamespaceLabel.from_str() result equals argument."""
+        label = provision.NamespaceLabel.from_str('ns:value')
+        self.assertEqual(label, 'ns:value')
+
+    def test_namespace_with_multiple_colons(self):
+        """Test NamespaceLabel.from_str() on argument with multiple colons."""
+        label = provision.NamespaceLabel.from_str('ns:value:value2')
+        self.assertEqual(label.namespace, 'ns')
+        self.assertEqual(label.value, 'value:value2')
+
+    def test_get_action(self):
+        """Test Label action property."""
+        action = provision.NamespaceLabel('cros-version', 'foo').action
+        self.assertEqual(action, provision.Action('cros-version', 'foo'))
+
+
+class CrosVersionLabelTestCase(unittest.TestCase):
+    """CrosVersionLabel test case."""
+
+    def test_value_is_cros_image(self):
+        """Test that value is CrosVersion type."""
+        label = provision.CrosVersionLabel('lumpy-release/R27-3773.0.0')
+        self.assertIsInstance(label.value, provision.CrosVersion)
+
+
+class CrosVersionTestCase(unittest.TestCase):
+    """CrosVersion test case."""
+
+    def test_cros_image_identity(self):
+        """Test that CrosVersion doesn't change string value."""
+        cros_image = provision.CrosVersion('lumpy-release/R27-3773.0.0')
+        self.assertEqual(cros_image, 'lumpy-release/R27-3773.0.0')
+
+    def test_cros_image_group(self):
+        """Test CrosVersion group property."""
+        cros_image = provision.CrosVersion('lumpy-release/R27-3773.0.0')
+        self.assertEqual(cros_image.group, 'lumpy-release')
+
+    def test_cros_image_group_na(self):
+        """Test invalid CrosVersion group property."""
+        cros_image = provision.CrosVersion('foo')
+        self.assertEqual(cros_image.group, cros_image.INVALID_STR)
+
+    def test_cros_image_milestone(self):
+        """Test CrosVersion milestone property."""
+        cros_image = provision.CrosVersion('lumpy-release/R27-3773.0.0')
+        self.assertEqual(cros_image.milestone, 'R27')
+
+    def test_cros_image_milestone_na(self):
+        """Test invalid CrosVersion milestone property."""
+        cros_image = provision.CrosVersion('foo')
+        self.assertEqual(cros_image.milestone, cros_image.INVALID_STR)
+
+    def test_cros_image_milestone_latest(self):
+        """Test that LATEST milestone isn't recognized."""
+        cros_image = provision.CrosVersion('lumpy-release/LATEST-3773.0.0')
+        self.assertEqual(cros_image.milestone, cros_image.INVALID_STR)
+
+    def test_cros_image_version(self):
+        """Test CrosVersion image property."""
+        cros_image = provision.CrosVersion('lumpy-release/R27-3773.0.0')
+        self.assertEqual(cros_image.version, '3773.0.0')
+
+    def test_cros_image_version_na(self):
+        """Test invalid CrosVersion version property."""
+        cros_image = provision.CrosVersion('foo')
+        self.assertEqual(cros_image.version, cros_image.INVALID_STR)
+
+    def test_cros_image_rc(self):
+        """Test CrosVersion rc property."""
+        cros_image = provision.CrosVersion('lumpy-release/R27-3773.0.0-rc2')
+        self.assertEqual(cros_image.rc, 'rc2')
+
+    def test_cros_image_rc_na(self):
+        """Test invalid CrosVersion rc property."""
+        cros_image = provision.CrosVersion('foo')
+        self.assertEqual(cros_image.rc, cros_image.INVALID_STR)
+
+    def test_cros_image_repr(self):
+        """Test CrosVersion.__repr__ works."""
+        cros_image = provision.CrosVersion('lumpy-release/R27-3773.0.0-rc2')
+        self.assertEqual(repr(cros_image),
+                         "CrosVersion('lumpy-release/R27-3773.0.0-rc2')")
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/server/hosts/cros_host.py b/server/hosts/cros_host.py
index 10e167c..0cd86c9 100644
--- a/server/hosts/cros_host.py
+++ b/server/hosts/cros_host.py
@@ -947,10 +947,10 @@
                         fwro_version and fwrw_version.
 
         """
-        fw_label = provision.fwrw_version_to_label(build)
+        fw_label = provision.FWRWVersionLabel(build)
         self._AFE.run('label_add_hosts', id=fw_label, hosts=[self.hostname])
         if not rw_only:
-            fw_label = provision.fwro_version_to_label(build)
+            fw_label = provision.FWROVersionLabel(build)
             self._AFE.run('label_add_hosts', id=fw_label, hosts=[self.hostname])