Update from https://crrev.com/333737

build/android/pylib/device/battery_utils.py

 # Copyright 2015 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 """Provides a variety of device interactions with power.
 """
 # pylint: disable=unused-argument
 
 import collections
 import contextlib
@@ skipping 271 matching lines @@
       device_errors.DeviceVersionError: If device is not L or higher.
     """
     def battery_updates_disabled():
       return self.GetCharging() is False
 
     if (self._device.build_version_sdk <
         constants.ANDROID_SDK_VERSION_CODES.LOLLIPOP):
       raise device_errors.DeviceVersionError('Device must be L or higher.')
 
     self._device.RunShellCommand(
-        ['dumpsys', 'battery', 'reset'], check_return=True)
+        ['dumpsys', 'battery', 'set', 'usb', '1'], check_return=True)
+    self._device.RunShellCommand(
+        ['dumpsys', 'battery', 'set', 'ac', '1'], check_return=True)
     self._device.RunShellCommand(
         ['dumpsys', 'batterystats', '--reset'], check_return=True)
     battery_data = self._device.RunShellCommand(
         ['dumpsys', 'batterystats', '--charged', '--checkin'],
-        check_return=True)
+        check_return=True, large_output=True)
     ROW_TYPE_INDEX = 3
     PWI_POWER_INDEX = 5
     for line in battery_data:
       l = line.split(',')
       if (len(l) > PWI_POWER_INDEX and l[ROW_TYPE_INDEX] == 'pwi'
           and l[PWI_POWER_INDEX] != 0):
         raise device_errors.CommandFailedError(
             'Non-zero pmi value found after reset.')
     self._device.RunShellCommand(['dumpsys', 'battery', 'set', 'ac', '0'],
                                  check_return=True)
@@ skipping 88 matching lines @@
       if temp is None:
         logging.warning('Unable to find current battery temperature.')
         temp = 0
       else:
         logging.info('Current battery temperature: %s', temp)
       return int(temp) <= target_temp
 
     logging.info('Waiting for the device to cool down to %s degrees.',
                  target_temp)
     timeout_retry.WaitFor(cool_device, wait_period=wait_period)