autotest: Update servo_state during DUT repair task

Updating servo_state during ServoHost.verify() and ServoHost.repair()

BUG=chromium:1045191,chromium:1012504
TEST=unittest

Change-Id: I794e90f0bbab554bd87fa89d8785eab9b0023f7c
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/autotest/+/2053074
Reviewed-by: Garry Wang <xianuowang@chromium.org>
Commit-Queue: Otabek Kasimov <otabek@google.com>
Tested-by: Otabek Kasimov <otabek@google.com>
Auto-Submit: Otabek Kasimov <otabek@google.com>

server/hosts/cros_host_unittest.py

diff --git a/server/hosts/cros_host_unittest.py b/server/hosts/cros_host_unittest.py
index 26b85b4..9dd4bc9 100755
--- a/server/hosts/cros_host_unittest.py
+++ b/server/hosts/cros_host_unittest.py
@@ -1,12 +1,14 @@
 #!/usr/bin/python2
 # pylint: disable=missing-docstring
 
+import mock
 import unittest
 
 import common
 
 from autotest_lib.server.hosts import cros_host
 from autotest_lib.server.hosts import servo_host
+from autotest_lib.server.hosts import host_info
 
 CROSSYSTEM_RESULT = '''
 fwb_tries              = 0                              # Fake comment
@@ -26,6 +28,9 @@
 CHROMEOS_RELEASE_UNIBUILD=1
 '''
 
+SERVO_STATE_PREFIX = servo_host.SERVO_STATE_LABEL_PREFIX
+
+
 class MockCmd(object):
     """Simple mock command with base command and results"""
 
@@ -104,6 +109,71 @@
         self.assertEqual(got, {servo_host.SERVO_HOST_ATTR: 'host'})
 
 
+class DictFilteringTestCase(unittest.TestCase):
+    """Test to verify servo_state was set-up as label in host_info_store"""
+
+    def create_host(self):
+        host = MockHost()
+        host.servo = None
+        host._servo_host = mock.Mock()
+        host._servo_host.get_servo.return_value = 'Not Empty'
+        host._servo_host.get_servo_state.return_value = 'SOME_STATE'
+        host.host_info_store = host_info.InMemoryHostInfoStore()
+        self.assertEqual(host.host_info_store.get().get_label_value(SERVO_STATE_PREFIX), '')
+        return host
+
+    def test_do_not_update_label_when_servo_host_is_not_inited(self):
+        host = self.create_host()
+        host._servo_host = None
+
+        host._update_servo_labels()
+        self.assertEqual(host.host_info_store.get().get_label_value(SERVO_STATE_PREFIX), '')
+
+    def test_do_not_update_label_when_servo_host_is_inited(self):
+        host = self.create_host()
+
+        host._update_servo_labels()
+        host._servo_host.get_servo_state.assert_called()
+        self.assertEqual(host.host_info_store.get().get_label_value(SERVO_STATE_PREFIX), 'SOME_STATE')
+
+    def test_repair_servo__update_servo_labels_after_repair_when_repair_is_fail(self):
+        host = self.create_host()
+        host._servo_host.repair.side_effect = Exception('Something bad')
+
+        try:
+            host.repair_servo()
+            self.assertEqual("Exception is", 'expecting to raise')
+        except:
+            pass
+        host._servo_host.get_servo_state.assert_called()
+        self.assertEqual(host.host_info_store.get().get_label_value(SERVO_STATE_PREFIX), 'SOME_STATE')
+
+    def test_repair_servo__update_servo_labels_after_repair_when_repair_is_not_fail(self):
+        host = self.create_host()
+        try:
+            host.repair_servo()
+        except:
+            self.assertEqual("Exception is not", 'expected')
+            pass
+        host._servo_host.get_servo_state.assert_called()
+        self.assertEqual(host.host_info_store.get().get_label_value(SERVO_STATE_PREFIX), 'SOME_STATE')
+
+    def test_set_servo_host_update_servo_state_when_host_exist(self):
+        host = self.create_host()
+        host._servo_host = mock.Mock()
+        host._servo_host.get_servo.return_value = 'Not Empty'
+        host._servo_host.get_servo_state.return_value = 'SOME_STATE'
+        self.assertEqual(host.host_info_store.get().get_label_value(SERVO_STATE_PREFIX), '')
+
+        try:
+            host.repair_servo()
+        except:
+            self.assertEqual("Exception is not", 'expected')
+            pass
+        host._servo_host.get_servo_state.assert_called()
+        self.assertEqual(host.host_info_store.get().get_label_value(SERVO_STATE_PREFIX), 'SOME_STATE')
+
+
 if __name__ == "__main__":
     unittest.main()