Add VBUS current check for type-C charging test

1. add poll_pd_state for plankton class. The state transition (between
sink and source) will take a short time. If we want to make sure test
checking in stable state, we should poll it.

2. add VBUS current check (same as the metric of factory test). Check
current is around 3A in charging 5V, and others current is over 0A.

3. make plankton object as local since it is only used inside run_once.

BUG=chrome-os-partner:33362
TEST=manual
  test_that --fast -b typec_host_board due_ip_address \
  firmware_TypeCCharging

Change-Id: I9a544b3a0e96a35df916069619895172612333e5
Reviewed-on: https://chromium-review.googlesource.com/262536
Reviewed-by: Wai-Hong Tam <waihong@chromium.org>
Commit-Queue: Pin-chih Lin <johnylin@chromium.org>
Tested-by: Pin-chih Lin <johnylin@chromium.org>
diff --git a/server/cros/servo/plankton.py b/server/cros/servo/plankton.py
index 482e42b..8154f34 100644
--- a/server/cros/servo/plankton.py
+++ b/server/cros/servo/plankton.py
@@ -7,6 +7,8 @@
 import time
 import xmlrpclib
 
+from autotest_lib.client.bin import utils
+
 
 class PlanktonError(Exception):
     pass
@@ -38,6 +40,11 @@
     VBUS_VOLTAGE_MV = 'vbus_voltage'
     VBUS_CURRENT_MA = 'vbus_current'
     VBUS_POWER_MW = 'vbus_power'
+    # USBC PD states.
+    USBC_PD_STATES = {
+        'sink': 'SNK_READY',
+        'source': 'SRC_READY'}
+    POLL_STATE_SECS = 2
 
 
     def __init__(self, args_dict=None):
@@ -125,3 +132,19 @@
             return 0
 
         raise PlanktonError('Invalid USBC role: %s' % usbc_role)
+
+
+    def poll_pd_state(self, state):
+        """Polls until Plankton pd goes to the specific state.
+
+        @param state: Specified pd state name.
+        """
+        if state not in self.USBC_PD_STATES:
+            raise PlanktonError('Invalid state name: %s' % state)
+        utils.poll_for_condition(
+            lambda: self.get('pd_state') == self.USBC_PD_STATES[state],
+            exception=utils.TimeoutError('Plankton not in %s state '
+                                         'after %s seconds.' %
+                                         (self.USBC_PD_STATES[state],
+                                          self.POLL_STATE_SECS)),
+            timeout=self.POLL_STATE_SECS)
diff --git a/server/site_tests/firmware_TypeCCharging/firmware_TypeCCharging.py b/server/site_tests/firmware_TypeCCharging/firmware_TypeCCharging.py
index 9dfc64c..8635018 100644
--- a/server/site_tests/firmware_TypeCCharging/firmware_TypeCCharging.py
+++ b/server/site_tests/firmware_TypeCCharging/firmware_TypeCCharging.py
@@ -17,28 +17,34 @@
 
     USBC_SINK_VOLTAGE = 5
     VBUS_TOLERANCE = 0.12
+    VBUS_5V_CURRENT_RANGE = (2, 3.4)
 
 
     def run_once(self, host, args_dict):
-        """Compares VBUS voltage with charging setting.
+        """Compares VBUS voltage and current with charging setting.
 
         When charging voltage == 0, Plankton will act as a power sink and draws
         5 volts from DUT. Other charging voltage should be seen on USB type C
         VBUS INA meter in a 12% range.
 
-        @raise TestFail: If VBUS voltage is not in range.
-        """
-        self._host = host
-        self._plankton = plankton.Plankton(args_dict)
+        When charging voltage == 5, Plankton INA current should be seen around
+        3 Amps (we set the range among 2 ~ 3.4 Amps just as in factory testing).
+        Other positive charging votage should not be less than 0 Amp.
 
-        for charging_voltage in self._plankton.get_charging_voltages():
-            self._plankton.charge(charging_voltage)
+        @raise TestFail: If VBUS voltage or current is not in range.
+        """
+        plankton_host = plankton.Plankton(args_dict)
+
+        for charging_voltage in plankton_host.get_charging_voltages():
+            plankton_host.charge(charging_voltage)
+            plankton_host.poll_pd_state(
+                    'source' if charging_voltage > 0 else 'sink')
             expected_vbus_voltage = float(
                     charging_voltage if charging_voltage > 0 else
                     self.USBC_SINK_VOLTAGE)
             tolerance = self.VBUS_TOLERANCE * expected_vbus_voltage
-            vbus_voltage = self._plankton.vbus_voltage
-            vbus_current = self._plankton.vbus_current
+            vbus_voltage = plankton_host.vbus_voltage
+            vbus_current = plankton_host.vbus_current
             logging.info('Charging %dV: VBUS V=%f I=%f', charging_voltage,
                          vbus_voltage, vbus_current)
 
@@ -49,3 +55,15 @@
 
             if charging_voltage == 0 and vbus_current > 0:
                 raise error.TestFail('Failed to consume power from DUT')
+
+            if charging_voltage > 0 and vbus_current <= 0:
+                raise error.Testfail(
+                        'VBUS current less than 0 in %d volt: %f' %
+                        (charging_voltage, vbus_current))
+
+            if (charging_voltage == 5 and
+                (vbus_current < self.VBUS_5V_CURRENT_RANGE[0] or
+                 vbus_current > self.VBUS_5V_CURRENT_RANGE[1])):
+                raise error.TestFail(
+                        'VBUS current out of range in 5 volt: %f %r' %
+                        (vbus_current, self.VBUS_5V_CURRENT_RANGE))