Move battery related functions to ChromeEC class

This patch moves update_battery_info() and get_battery*() function
family to ChromeEC class, so these functions can be used in other tests
too.

BUG=b:161775827
TEST=Run firmware_ECCharging test, make sure test works

Change-Id: I8701c7d0eb2768c2aade275beaa3b64206de3630
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/autotest/+/2753987
Tested-by: Patryk Duda <pdk@semihalf.com>
Reviewed-by: Wai-Hong Tam <waihong@google.com>
Commit-Queue: Wai-Hong Tam <waihong@google.com>
diff --git a/server/cros/servo/chrome_ec.py b/server/cros/servo/chrome_ec.py
index 9d7f524..3eb0b4c 100644
--- a/server/cros/servo/chrome_ec.py
+++ b/server/cros/servo/chrome_ec.py
@@ -249,6 +249,9 @@
     This class is to abstract these interfaces.
     """
 
+    # The dict to cache the battery information
+    BATTERY_INFO = {}
+
     def __init__(self, servo, name="ec_uart"):
         super(ChromeEC, self).__init__(servo, name)
 
@@ -440,6 +443,110 @@
 
         return ((1 << (feat_id - feat_start)) & feat_bitmap) != 0
 
+    def update_battery_info(self):
+        """Get the battery info we care for this test."""
+        # The battery parameters we care for this test. The order must match
+        # the output of EC battery command.
+        battery_params = [
+                'V', 'V-desired', 'I', 'I-desired', 'Charging', 'Remaining'
+        ]
+        regex_str_list = []
+
+        for p in battery_params:
+            if p == 'Remaining':
+                regex_str_list.append(p + ':\s+(\d+)\s+')
+            elif p == 'Charging':
+                regex_str_list.append(p + r':\s+(Not Allowed|Allowed)\s+')
+            else:
+                regex_str_list.append(p +
+                                      r':\s+0x[0-9a-f]*\s+=\s+([0-9-]+)\s+')
+
+        # For unknown reasons, servod doesn't always capture the ec
+        # command output. It doesn't happen often, but retry if it does.
+        retries = 3
+        while retries > 0:
+            retries -= 1
+            try:
+                battery_regex_match = self.send_command_get_output(
+                        'battery', regex_str_list)
+                break
+            except (servo.UnresponsiveConsoleError,
+                    servo.ResponsiveConsoleError) as e:
+                if retries <= 0:
+                    raise
+                logging.warning('Failed to get battery status. %s', e)
+        else:
+            battery_regex_match = self.send_command_get_output(
+                    'battery', regex_str_list)
+
+        for i in range(len(battery_params)):
+            if battery_params[i] == 'Charging':
+                self.BATTERY_INFO[
+                        battery_params[i]] = battery_regex_match[i][1]
+            else:
+                self.BATTERY_INFO[battery_params[i]] = int(
+                        battery_regex_match[i][1])
+        logging.debug('Battery info: %s', self.BATTERY_INFO)
+
+    def get_battery_desired_voltage(self, print_result=True):
+        """Get battery desired voltage value."""
+        if not self.BATTERY_INFO:
+            self.update_battery_info()
+        if print_result:
+            logging.info('Battery desired voltage = %d mV',
+                         self.BATTERY_INFO['V-desired'])
+        return self.BATTERY_INFO['V-desired']
+
+    def get_battery_desired_current(self, print_result=True):
+        """Get battery desired current value."""
+        if not self.BATTERY_INFO:
+            self.update_battery_info()
+        if print_result:
+            logging.info('Battery desired current = %d mA',
+                         self.BATTERY_INFO['I-desired'])
+        return self.BATTERY_INFO['I-desired']
+
+    def get_battery_actual_voltage(self, print_result=True):
+        """Get the actual voltage from charger to battery."""
+        if not self.BATTERY_INFO:
+            self.update_battery_info()
+        if print_result:
+            logging.info('Battery actual voltage = %d mV',
+                         self.BATTERY_INFO['V'])
+        return self.BATTERY_INFO['V']
+
+    def get_battery_actual_current(self, print_result=True):
+        """Get the actual current from charger to battery."""
+        if not self.BATTERY_INFO:
+            self.update_battery_info()
+        if print_result:
+            logging.info('Battery actual current = %d mA',
+                         self.BATTERY_INFO['I'])
+        return self.BATTERY_INFO['I']
+
+    def get_battery_remaining(self, print_result=True):
+        """Get battery charge remaining in mAh."""
+        if not self.BATTERY_INFO:
+            self.update_battery_info()
+        if print_result:
+            logging.info("Battery charge remaining = %d mAh",
+                         self.BATTERY_INFO['Remaining'])
+        return self.BATTERY_INFO['Remaining']
+
+    def get_battery_charging_allowed(self, print_result=True):
+        """Get the battery charging state.
+
+        Returns True if charging is allowed.
+        """
+        if not self.BATTERY_INFO:
+            self.update_battery_info()
+        if print_result:
+            logging.info("Battery charging = %s",
+                         self.BATTERY_INFO['Charging'])
+        if self.BATTERY_INFO['Charging'] == 'Allowed':
+            return True
+        return False
+
 
 class ChromeUSBPD(ChromeEC):
     """Manages control of a Chrome USBPD.
diff --git a/server/site_tests/firmware_ECCharging/firmware_ECCharging.py b/server/site_tests/firmware_ECCharging/firmware_ECCharging.py
index 8dec8f6..cab6257 100644
--- a/server/site_tests/firmware_ECCharging/firmware_ECCharging.py
+++ b/server/site_tests/firmware_ECCharging/firmware_ECCharging.py
@@ -9,7 +9,6 @@
 from autotest_lib.server.cros.faft.firmware_test import FirmwareTest
 from autotest_lib.server.cros.servo import servo
 
-
 class firmware_ECCharging(FirmwareTest):
     """
     Servo based EC charging control test.
@@ -42,9 +41,6 @@
     # This should be >= BEGIN_CHARGING_TIMEOUT
     EXTRA_DISCHARGE_TIME = BEGIN_CHARGING_TIMEOUT + 30
 
-    # The dict to cache the battery information
-    BATTERY_INFO = {}
-
     def initialize(self, host, cmdline_args):
         super(firmware_ECCharging, self).initialize(host, cmdline_args)
         # Don't bother if there is no Chrome EC.
@@ -62,101 +58,6 @@
             logging.error("Caught exception: %s", str(e))
         super(firmware_ECCharging, self).cleanup()
 
-    def _update_battery_info(self):
-        """Get the battery info we care for this test."""
-        # The battery parameters we care for this test. The order must match
-        # the output of EC battery command.
-        battery_params = [
-                'V', 'V-desired', 'I', 'I-desired', 'Charging', 'Remaining'
-        ]
-        regex_str_list = []
-
-        for p in battery_params:
-            if p == 'Remaining':
-                regex_str_list.append(p + ':\s+(\d+)\s+')
-            elif p == 'Charging':
-                regex_str_list.append(p + r':\s+(Not Allowed|Allowed)\s+')
-            else:
-                regex_str_list.append(p +
-                                      r':\s+0x[0-9a-f]*\s+=\s+([0-9-]+)\s+')
-
-        # For unknown reasons, servod doesn't always capture the ec
-        # command output. It doesn't happen often, but retry if it does.
-        retries = 3
-        while retries > 0:
-            retries -= 1
-            try:
-                battery_regex_match = self.ec.send_command_get_output(
-                        'battery', regex_str_list)
-                break
-            except (servo.UnresponsiveConsoleError,
-                    servo.ResponsiveConsoleError) as e:
-                if retries <= 0:
-                    raise
-                logging.warning('Failed to get battery status. %s', e)
-        else:
-            battery_regex_match = self.ec.send_command_get_output(
-                    'battery', regex_str_list)
-
-        for i in range(len(battery_params)):
-            if battery_params[i] == 'Charging':
-                self.BATTERY_INFO[
-                        battery_params[i]] = battery_regex_match[i][1]
-            else:
-                self.BATTERY_INFO[battery_params[i]] = int(
-                        battery_regex_match[i][1])
-        logging.debug('Battery info: %s', self.BATTERY_INFO)
-
-    def _get_battery_desired_voltage(self):
-        """Get battery desired voltage value."""
-        if not self.BATTERY_INFO:
-            self._update_battery_info()
-        logging.info('Battery desired voltage = %d mV',
-                     self.BATTERY_INFO['V-desired'])
-        return self.BATTERY_INFO['V-desired']
-
-    def _get_battery_desired_current(self):
-        """Get battery desired current value."""
-        if not self.BATTERY_INFO:
-            self._update_battery_info()
-        logging.info('Battery desired current = %d mA',
-                     self.BATTERY_INFO['I-desired'])
-        return self.BATTERY_INFO['I-desired']
-
-    def _get_battery_actual_voltage(self):
-        """Get the actual voltage from charger to battery."""
-        if not self.BATTERY_INFO:
-            self._update_battery_info()
-        logging.info('Battery actual voltage = %d mV', self.BATTERY_INFO['V'])
-        return self.BATTERY_INFO['V']
-
-    def _get_battery_actual_current(self):
-        """Get the actual current from charger to battery."""
-        if not self.BATTERY_INFO:
-            self._update_battery_info()
-        logging.info('Battery actual current = %d mA', self.BATTERY_INFO['I'])
-        return self.BATTERY_INFO['I']
-
-    def _get_battery_remaining(self):
-        """Get battery charge remaining in mAh."""
-        if not self.BATTERY_INFO:
-            self._update_battery_info()
-        logging.info("Battery charge remaining = %d mAh",
-                     self.BATTERY_INFO['Remaining'])
-        return self.BATTERY_INFO['Remaining']
-
-    def _get_battery_charging_allowed(self):
-        """Get the battery charging state.
-
-        Returns True if charging is allowed.
-        """
-        if not self.BATTERY_INFO:
-            self._update_battery_info()
-        logging.info("Battery charging = %s", self.BATTERY_INFO['Charging'])
-        if self.BATTERY_INFO['Charging'] == 'Allowed':
-            return True
-        return False
-
     def _get_charger_target_voltage(self):
         """Get target charging voltage set in charger."""
         voltage = int(
@@ -175,7 +76,7 @@
 
     def _get_trickle_charging(self):
         """Check if we are trickle charging battery."""
-        return (self._get_battery_desired_current() <
+        return (self.ec.get_battery_desired_current() <
                 self.TRICKLE_CHARGE_THRESHOLD)
 
     def _check_target_value(self):
@@ -185,21 +86,21 @@
           error.TestFail: Raised when check fails.
         """
         if (self._get_charger_target_voltage() >=
-                1.05 * self._get_battery_desired_voltage()):
+                    1.05 * self.ec.get_battery_desired_voltage()):
             raise error.TestFail(
                     "Charger target voltage is too high. %d/%d=%f" %
                     (self._get_charger_target_voltage(),
-                     self._get_battery_desired_voltage(),
+                     self.ec.get_battery_desired_voltage(),
                      float(self._get_charger_target_voltage()) /
-                     self._get_battery_desired_voltage()))
+                     self.ec.get_battery_desired_voltage()))
         if (self._get_charger_target_current() >=
-                1.05 * self._get_battery_desired_current()):
+                    1.05 * self.ec.get_battery_desired_current()):
             raise error.TestFail(
                     "Charger target current is too high. %d/%d=%f" %
                     (self._get_charger_target_current(),
-                     self._get_battery_desired_current(),
+                     self.ec.get_battery_desired_current(),
                      float(self._get_charger_target_current()) /
-                     self._get_battery_desired_current()))
+                     self.ec.get_battery_desired_current()))
 
     def _check_actual_value(self):
         """Check actual voltage/current values are correct.
@@ -207,21 +108,21 @@
         Raise:
           error.TestFail: Raised when check fails.
         """
-        if (self._get_battery_actual_voltage() >=
-                1.05 * self._get_charger_target_voltage()):
+        if (self.ec.get_battery_actual_voltage() >=
+                    1.05 * self._get_charger_target_voltage()):
             raise error.TestFail(
                     "Battery actual voltage is too high. %d/%d=%f" %
-                    (self._get_battery_actual_voltage(),
+                    (self.ec.get_battery_actual_voltage(),
                      self._get_charger_target_voltage(),
-                     float(self._get_battery_actual_voltage()) /
+                     float(self.ec.get_battery_actual_voltage()) /
                      self._get_charger_target_voltage()))
-        if (self._get_battery_actual_current() >=
-                1.05 * self._get_charger_target_current()):
+        if (self.ec.get_battery_actual_current() >=
+                    1.05 * self._get_charger_target_current()):
             raise error.TestFail(
                     "Battery actual current is too high. %d/%d=%f" %
-                    (self._get_battery_actual_current(),
+                    (self.ec.get_battery_actual_current(),
                      self._get_charger_target_current(),
-                     float(self._get_battery_actual_current()) /
+                     float(self.ec.get_battery_actual_current()) /
                      self._get_charger_target_current()))
 
     def _check_if_discharge_on_ac(self):
@@ -279,7 +180,7 @@
         # Verify AC is on and charge control is normal.
         if self._check_battery_discharging():
             raise error.TestFail("Fail to plug AC and enable charging.")
-        self._update_battery_info()
+        self.ec.update_battery_info()
 
     def _consume_battery(self, deadline):
         """Perform battery intensive operation to make the battery discharge faster."""
@@ -300,11 +201,11 @@
             # Wait until DISCHARGE_TIMEOUT or charging allowed
             deadline = time.time() + self.DISCHARGE_TIMEOUT
             while time.time() < deadline:
-                self._update_battery_info()
-                if self._get_battery_charging_allowed():
+                self.ec.update_battery_info()
+                if self.ec.get_battery_charging_allowed():
                     break
                 logging.info("Wait for the battery to discharge (%d mAh).",
-                             self._get_battery_remaining())
+                             self.ec.get_battery_remaining())
                 self._consume_battery(deadline)
             else:
                 raise error.TestFail(
@@ -314,10 +215,10 @@
             # Wait another EXTRA_DISCHARGE_TIME just to be sure
             deadline = time.time() + self.EXTRA_DISCHARGE_TIME
             while time.time() < deadline:
-                self._update_battery_info()
+                self.ec.update_battery_info()
                 logging.info(
                         "Wait for the battery to discharge even more (%d mAh).",
-                        self._get_battery_remaining())
+                        self.ec.get_battery_remaining())
                 self._consume_battery(deadline)
         finally:
             self._set_battery_normal()
@@ -326,13 +227,13 @@
         # battery to actually start charging.
         deadline = time.time() + self.BEGIN_CHARGING_TIMEOUT
         while time.time() < deadline:
-            self._update_battery_info()
-            if self._get_battery_actual_current() >= 0:
+            self.ec.update_battery_info()
+            if self.ec.get_battery_actual_current() >= 0:
                 break
             logging.info(
                     'Battery actual current (%d) too low, wait a bit. (%d mAh)',
-                    self._get_battery_actual_current(),
-                    self._get_battery_remaining())
+                    self.ec.get_battery_actual_current(),
+                    self.ec.get_battery_remaining())
             time.sleep(self.BEGIN_CHARGING_RETRY_TIME)
 
     def run_once(self):
@@ -341,13 +242,13 @@
         if not self.check_ec_capability(['battery', 'charging']):
             raise error.TestNAError(
                     "Nothing needs to be tested on this device")
-        if not self._get_battery_charging_allowed(
-        ) or self._get_battery_actual_current() < 0:
+        if not self.ec.get_battery_charging_allowed(
+        ) or self.ec.get_battery_actual_current() < 0:
             logging.info(
                     "Battery is full or discharging. Forcing battery discharge to test charging."
             )
             self._discharge_below_100()
-            if not self._get_battery_charging_allowed():
+            if not self.ec.get_battery_charging_allowed():
                 raise error.TestFail(
                         'Battery reports charging is not allowed, even after discharging.'
                 )
@@ -357,7 +258,7 @@
         if self._get_trickle_charging():
             raise error.TestNAError(
                     "Trickling charging battery. Unable to test.")
-        if self._get_battery_actual_current() < 0:
+        if self.ec.get_battery_actual_current() < 0:
             raise error.TestFail(
                     "The device is not charging. Is the test run with AC plugged?"
             )