autotest: DUT label update refactoring for repair and deploy task am: 871a2c037e am: 0050caf4d2 am: 271d313dd8 am: 94e7daaddc

Change-Id: Ifa8de32460bb37f748c90ab4ae23f84c92aaab99
diff --git a/server/control_segments/repair b/server/control_segments/repair
index 8aace51..35cd9b9 100644
--- a/server/control_segments/repair
+++ b/server/control_segments/repair
@@ -43,7 +43,9 @@
         logging.debug('Repair with labels list %s', labels_list)
 
         try:
-            target.labels.update_labels(target, keep_pool=True)
+            target.labels.update_labels(target,
+                                        task_name='repair',
+                                        keep_pool=True)
         except Exception:
             logging.exception('Exception while updating labels.')
     except Exception:
diff --git a/server/hosts/base_label.py b/server/hosts/base_label.py
index 3cfe7f6..ec4cc6f 100644
--- a/server/hosts/base_label.py
+++ b/server/hosts/base_label.py
@@ -102,6 +102,25 @@
         return prefix_labels, full_labels
 
 
+    def update_for_task(self, task_name):
+        """
+        This method helps to check which labels need to be updated.
+        State config labels are updated only for repair task.
+        Lab config labels are updated only for deploy task.
+        All labels are updated for any task.
+
+        It is the responsibility of the subclass to override this method
+        to differentiate itself as a state config label or a lab config label
+        and return the appropriate boolean value.
+
+        If the subclass doesn't override this method then that label will
+        always be updated for any type of task.
+
+        @returns True if labels should be updated for the task with given name
+        """
+        return True
+
+
 class StringLabel(BaseLabel):
     """
     This class represents a string label that is dynamically generated.
@@ -174,12 +193,14 @@
 class LabelRetriever(object):
     """This class will assist in retrieving/updating the host labels."""
 
-    def _populate_known_labels(self, label_list):
+    def _populate_known_labels(self, label_list, task_name):
         """Create a list of known labels that is created through this class."""
         for label_instance in label_list:
-            prefixed_labels, full_labels = label_instance.get_all_labels()
-            self.label_prefix_names.update(prefixed_labels)
-            self.label_full_names.update(full_labels)
+            # populate only the labels that need to be updated for this task.
+            if label_instance.update_for_task(task_name):
+                prefixed_labels, full_labels = label_instance.get_all_labels()
+                self.label_prefix_names.update(prefixed_labels)
+                self.label_full_names.update(full_labels)
 
 
     def __init__(self, label_list):
@@ -207,6 +228,28 @@
         return labels
 
 
+    def get_labels_for_update(self, host, task_name):
+        """
+        Retrieve the labels for the host which needs to be updated.
+
+        @param host: The host to get the labels for updating.
+        @param task_name: task name(repair/deploy) for the operation.
+
+        @returns labels to be updated
+        """
+        labels = []
+        for label in self._labels:
+            logging.info('checking label update %s', label.__class__.__name__)
+            try:
+                # get only the labels which need to be updated for this task.
+                if label.update_for_task(task_name):
+                    labels.extend(label.get(host))
+            except Exception:
+                logging.exception('error getting label %s.',
+                                  label.__class__.__name__)
+        return labels
+
+
     def _is_known_label(self, label):
         """
         Checks if the label is a label known to the label detection framework.
@@ -253,7 +296,7 @@
         host.host_info_store.commit(new_info)
 
 
-    def update_labels(self, host, keep_pool=False):
+    def update_labels(self, host, task_name='', keep_pool=False):
         """
         Retrieve the labels from the host and update if needed.
 
@@ -261,12 +304,12 @@
         """
         # If we haven't yet grabbed our list of known labels, do so now.
         if not self.label_full_names and not self.label_prefix_names:
-            self._populate_known_labels(self._labels)
+            self._populate_known_labels(self._labels, task_name)
 
         # Label detection hits the DUT so it can be slow. Do it before reading
         # old labels from HostInfoStore to minimize the time between read and
         # commit of the HostInfo.
-        new_labels = self.get_labels(host)
+        new_labels = self.get_labels_for_update(host, task_name)
         old_info = host.host_info_store.get()
         self._carry_over_unknown_labels(old_info.labels, new_labels)
         new_info = host_info.HostInfo(
diff --git a/server/hosts/base_label_unittest.py b/server/hosts/base_label_unittest.py
index 9710bc6..2a6db21 100755
--- a/server/hosts/base_label_unittest.py
+++ b/server/hosts/base_label_unittest.py
@@ -107,6 +107,9 @@
         self.assertEqual(prefix_tbl, set())
         self.assertEqual(prefix_tbls, set())
 
+    def test_update_for_task(self):
+        self.assertTrue(self.test_base_label.update_for_task(''))
+
 
 class StringPrefixLabelUnittests(unittest.TestCase):
     """Unittest for testing base_label.StringPrefixLabel."""
@@ -136,7 +139,7 @@
     def setUp(self):
         label_list = [TestStringPrefixLabel(), TestBaseLabel()]
         self.retriever = base_label.LabelRetriever(label_list)
-        self.retriever._populate_known_labels(label_list)
+        self.retriever._populate_known_labels(label_list, '')
 
 
     def test_populate_known_labels(self):
diff --git a/server/hosts/cros_label.py b/server/hosts/cros_label.py
index 065ce9e..17038db 100644
--- a/server/hosts/cros_label.py
+++ b/server/hosts/cros_label.py
@@ -27,6 +27,10 @@
 # fallback values if we can't contact the HWID server
 HWID_LABELS_FALLBACK = ['sku', 'phase', 'touchscreen', 'touchpad', 'variant', 'stylus']
 
+# Repair and Deploy taskName
+REPAIR_TASK_NAME = 'repair'
+DEPLOY_TASK_NAME = 'deploy'
+
 
 def _parse_lsb_output(host):
     """Parses the LSB output and returns key data points for labeling.
@@ -116,6 +120,11 @@
 
         return []
 
+    def update_for_task(self, task_name):
+        # This label is stored in the lab config, so only deploy tasks update it
+        # or when no task name is mentioned.
+        return task_name in (DEPLOY_TASK_NAME, '')
+
 
 class BrandCodeLabel(base_label.StringPrefixLabel):
     """Determine the correct brand_code (aka RLZ-code) for the device."""
@@ -305,6 +314,10 @@
             pass
         return has_chameleon
 
+    def update_for_task(self, task_name):
+        # This label is stored in the state config, so only repair tasks update
+        # it or when no task name is mentioned.
+        return task_name in (REPAIR_TASK_NAME, '')
 
 
 class ChameleonConnectionLabel(base_label.StringPrefixLabel):
@@ -319,6 +332,11 @@
     def generate_labels(self, host):
         return [chameleon.get_label() for chameleon in host.chameleon_list]
 
+    def update_for_task(self, task_name):
+        # This label is stored in the lab config, so only deploy tasks update it
+        # or when no task name is mentioned.
+        return task_name in (DEPLOY_TASK_NAME, '')
+
 
 class ChameleonPeripheralsLabel(base_label.StringPrefixLabel):
     """Return the Chameleon peripherals labels.
@@ -386,7 +404,10 @@
         logging.info('Bluetooth labels are %s', labels_list)
         return labels_list
 
-
+    def update_for_task(self, task_name):
+        # This label is stored in the lab config, so only deploy tasks update it
+        # or when no task name is mentioned.
+        return task_name in (DEPLOY_TASK_NAME, '')
 
 
 class AudioLoopbackDongleLabel(base_label.BaseLabel):
@@ -425,6 +446,11 @@
             return True
         return False
 
+    def update_for_task(self, task_name):
+        # This label is stored in the state config, so only repair tasks update
+        # it or when no task name is mentioned.
+        return task_name in (REPAIR_TASK_NAME, '')
+
 
 class PowerSupplyLabel(base_label.StringPrefixLabel):
     """
@@ -611,6 +637,11 @@
         return (servo_host_hostname is not None
                 and servo_host.servo_host_is_up(servo_host_hostname))
 
+    def update_for_task(self, task_name):
+        # This label is stored in the state config, so only repair tasks update
+        # it or when no task name is mentioned.
+        return task_name in (REPAIR_TASK_NAME, '')
+
 
 class ArcLabel(base_label.BaseLabel):
     """Label indicates if host has ARC support."""
diff --git a/server/hosts/cros_label_unittest.py b/server/hosts/cros_label_unittest.py
index f837e3e..1cd9c99 100755
--- a/server/hosts/cros_label_unittest.py
+++ b/server/hosts/cros_label_unittest.py
@@ -20,6 +20,10 @@
 from autotest_lib.server.hosts.cros_label import Cr50RWVersionLabel
 from autotest_lib.server.hosts.cros_label import DeviceSkuLabel
 from autotest_lib.server.hosts.cros_label import ModelLabel
+from autotest_lib.server.hosts.cros_label import AudioLoopbackDongleLabel
+from autotest_lib.server.hosts.cros_label import ChameleonConnectionLabel
+from autotest_lib.server.hosts.cros_label import ChameleonLabel
+from autotest_lib.server.hosts.cros_label import ChameleonPeripheralsLabel
 from autotest_lib.server.hosts.cros_label import ServoLabel
 from autotest_lib.server.hosts import host_info
 
@@ -239,6 +243,11 @@
         host = MockHost(['device-sku:48'])
         self.assertEqual(DeviceSkuLabel().generate_labels(host), ['48'])
 
+    def test_update_for_task(self):
+        self.assertTrue(DeviceSkuLabel().update_for_task(''))
+        self.assertFalse(DeviceSkuLabel().update_for_task('repair'))
+        self.assertTrue(DeviceSkuLabel().update_for_task('deploy'))
+
 
 class BrandCodeLabelTests(unittest.TestCase):
     """Unit tests for DeviceSkuLabel"""
@@ -440,6 +449,39 @@
           self.assertEqual(prefix_labels, set(['servo_state']))
           self.assertEqual(full_labels, set(['servo']))
 
+    def test_update_for_task(self):
+        self.assertTrue(ServoLabel().update_for_task(''))
+        self.assertTrue(ServoLabel().update_for_task('repair'))
+        self.assertFalse(ServoLabel().update_for_task('deploy'))
+
+
+class AudioLoopbackDongleLabelTests(unittest.TestCase):
+    def test_update_for_task(self):
+        self.assertTrue(AudioLoopbackDongleLabel().update_for_task(''))
+        self.assertTrue(AudioLoopbackDongleLabel().update_for_task('repair'))
+        self.assertFalse(AudioLoopbackDongleLabel().update_for_task('deploy'))
+
+
+class ChameleonConnectionLabelTests(unittest.TestCase):
+    def test_update_for_task(self):
+        self.assertTrue(ChameleonConnectionLabel().update_for_task(''))
+        self.assertFalse(ChameleonConnectionLabel().update_for_task('repair'))
+        self.assertTrue(ChameleonConnectionLabel().update_for_task('deploy'))
+
+
+class ChameleonLabelTests(unittest.TestCase):
+    def test_update_for_task(self):
+        self.assertTrue(ChameleonLabel().update_for_task(''))
+        self.assertTrue(ChameleonLabel().update_for_task('repair'))
+        self.assertFalse(ChameleonLabel().update_for_task('deploy'))
+
+
+class ChameleonPeripheralsLabelTests(unittest.TestCase):
+    def test_update_for_task(self):
+        self.assertTrue(ChameleonPeripheralsLabel().update_for_task(''))
+        self.assertFalse(ChameleonPeripheralsLabel().update_for_task('repair'))
+        self.assertTrue(ChameleonPeripheralsLabel().update_for_task('deploy'))
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/site_utils/deployment/prepare/main.py b/site_utils/deployment/prepare/main.py
index 836101d..a5cb687 100755
--- a/site_utils/deployment/prepare/main.py
+++ b/site_utils/deployment/prepare/main.py
@@ -94,7 +94,7 @@
 
     if 'update-label' in opts.actions:
       try:
-        host.labels.update_labels(host)
+        host.labels.update_labels(host, task_name='deploy')
       except Exception as err:
         logging.error("fail to update label: %s", err)
         return RETURN_CODES.UPDATE_LABEL_FAILURE