autotest: Add verifier for battery audit

BUG=b:163572001
TEST=run local repair again bad and good battery
bad battery chromeos2-row5-rack7-host17
good battery chromeos6-row12-rack5-host10

Change-Id: Ibc9e00757e532ecbb86c0c23c8ae6288699c3e08
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/autotest/+/2738887
Commit-Queue: Otabek Kasimov <otabek@google.com>
Tested-by: Otabek Kasimov <otabek@google.com>
Reviewed-by: Garry Wang <xianuowang@chromium.org>
diff --git a/server/hosts/cros_repair.py b/server/hosts/cros_repair.py
index ddd9325..2b967f8 100644
--- a/server/hosts/cros_repair.py
+++ b/server/hosts/cros_repair.py
@@ -32,6 +32,7 @@
 from autotest_lib.server.hosts import repair_utils
 from autotest_lib.site_utils.admin_audit import verifiers as audit_verify
 from autotest_lib.site_utils.admin_audit import constants as audit_const
+from autotest_lib.site_utils.admin_audit import battery_validator
 from six.moves import range
 
 try:
@@ -892,6 +893,43 @@
         return 'Ensure DUT storage SMART information is in good state.'
 
 
+class AuditBattery(hosts.Verifier):
+    """Verify that battery on DUT is good to use.
+
+    Check if DUT drive is providing good SMART stats which not showing any
+    issues on it. The verifier can mark DUT for replacement if SMART stats
+    show outworn data.
+    """
+
+    @timeout_util.TimeoutDecorator(cros_constants.VERIFY_TIMEOUT_SEC)
+    def verify(self, host):
+        # pylint: disable=missing-docstring
+        state = None
+        try:
+            state = self._get_validator(host).validate()
+        except Exception as e:
+            # We do not want stop main process if it fail.
+            logging.debug('(Not critical) %s', e)
+        if not state:
+            raise hosts.AutoservNonCriticalVerifyError(
+                    'DUT battery did not detected or state cannot extracted.')
+        if state == audit_const.HW_STATE_NEED_REPLACEMENT:
+            logging.info('Detected issue with storage on the DUT.')
+            host.set_device_needs_replacement()
+
+    def _is_applicable(self, host):
+        return self._get_validator(host).is_battery_expected()
+
+    def _get_validator(self, host):
+        if not getattr(self, '_validator', None):
+            self._validator = battery_validator.BatteryValidator(host)
+        return self._validator
+
+    @property
+    def description(self):
+        return 'Ensure DUT battery is in good state.'
+
+
 class ServoKeyboardMapVerifier(hosts.Verifier):
     """Not critical verify to flash servo keyboard for the host.
 
@@ -1537,6 +1575,7 @@
     return (
             (StopStartUIVerifier, 'stop_start_ui', ('ssh', )),
             (DUTStorageVerifier, 'storage', ('ssh', )),
+            (AuditBattery, 'audit_battery', ()),
             (GscToolPresentVerifier, 'dut_gsctool', ('ssh', )),
             (ServoKeyboardMapVerifier, 'dut_servo_keyboard', ('ssh', )),
             (ServoMacAddressVerifier, 'dut_servo_macaddr', ('ssh', )),
diff --git a/server/hosts/cros_repair_unittest.py b/server/hosts/cros_repair_unittest.py
index 197f773..1c5a23d 100755
--- a/server/hosts/cros_repair_unittest.py
+++ b/server/hosts/cros_repair_unittest.py
@@ -38,6 +38,7 @@
          ('ssh', )),
         (cros_repair.StopStartUIVerifier, 'stop_start_ui', ('ssh', )),
         (cros_repair.DUTStorageVerifier, 'storage', ('ssh', )),
+        (cros_repair.AuditBattery, 'audit_battery', ()),
         (cros_repair.GscToolPresentVerifier, 'dut_gsctool', ('ssh', )),
         (cros_repair.ServoKeyboardMapVerifier, 'dut_servo_keyboard',
          ('ssh', )),