Example of how to figure out which BattOr maps to which port [DO NOT SUBMIT]

systrace/systrace/agents/battor_trace_agent.py

Index: systrace/systrace/agents/battor_trace_agent.py
diff --git a/systrace/systrace/agents/battor_trace_agent.py b/systrace/systrace/agents/battor_trace_agent.py
new file mode 100644
index 0000000000000000000000000000000000000000..2e1f6f959f81d1c38b850ba46a07e3b332694d03
--- /dev/null
+++ b/systrace/systrace/agents/battor_trace_agent.py
@@ -0,0 +1,180 @@
+# Copyright (c) 2016 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.
+
+import subprocess
+import threading
+
+from systrace.tracing_agents import TracingAgent, TraceResults, usb_hubs
+
+from devil.utils import find_usb_devices
+from devil.android import battery_utils, device_utils
+from py_trace_event import trace_time
+
+
+def try_create_agent(options, categories, disable=False):
+  if disable:
+    return False
+  if options.from_file is not None:
+    return False
+  return BattorTraceAgent(options, categories)
+
+class BattorError(Exception):
+  pass
+
+class BattorTraceAgent(TracingAgent):
+  # Class representing tracing agent that gets data from a BattOr.
+  def __init__(self, options, categories):
+    """Initialize a BattOr tracing agent.
+
+    Args:
+      options: The command-line options.
+      categories: The trace categories to capture.
+    """
+    super(BattorTraceAgent, self).__init__()
+    self._trace_result = ''
+    self._devstring = ''
+    self._options = options
+    self._categories = categories
+    self._collection_process = None
+    self._battor_process = None
+    self._recording_error = None
+    self._device_utils = device_utils.DeviceUtils(self._options.device_serial)
+    self._battery_utils = battery_utils.BatteryUtils(self._device_utils)
+    self._devstring = self._GetDeviceString()
+
+  def _GetDeviceString(self):
+    """Gets the device string to communicate with device.
+
+    Returns:
+      Device string used to communicate with device.
+
+    Raises:
+      ValueError: If serial number is not given.
+      BattorError: If BattOr not found or unexpected USB topology.
+    """
+    # If there's only one BattOr connected to the system, just use that one.
+    # This allows for use on, e.g., a developer's workstation with no hubs.
+    serial = self._options.device_serial
+    devtree = find_usb_devices.GetBusNumberToDeviceTreeMap(fast=True)
+    all_battors = find_usb_devices.GetBattorList(devtree)
+    hub_types = [usb_hubs.GetHubType(x) for x in self._options.hub_types]
+    if len(all_battors) == 1:
+      return '/dev/' + all_battors[0]
+    if serial:
+      p2serial = find_usb_devices.GetAllPhysicalPortToSerialMaps(
+          hub_types, device_tree_map=devtree)
+      p2tty = find_usb_devices.GetAllPhysicalPortToTTYMaps(
+          hub_types, device_tree_map=devtree)
+
+      # get the port number of this device
+      port_num = -1
+      for hub in p2serial:
+        for (port, s) in hub.iteritems():
+          if serial == s:
+            if port_num != -1:
+              raise BattorError('Two devices with same serial number.')
+            else:
+              port_num = port
+      if port_num == -1:
+        raise BattorError('Device with given serial number not found.')
+
+      # get the tty for this port number
+      tty_string = None
+      for hub in p2tty:
+        if port_num in hub:
+          if tty_string:
+            raise BattorError('Two TTY devices with matching port number.')
+          else:
+            tty_string = hub[port_num]
+      if not tty_string:
+        raise BattorError('No BattOr detected on corresponding port.')
+      if find_usb_devices.IsBattor(tty_string, devtree):
+        return '/dev/' + tty_string
+      else:
+        raise BattorError('Device connected to matching port is not BattOr.')
+    else:
+      raise BattorError('Two or more BattOrs connected, no serial provided')
+
+  def StartAgentTracing(self):
+    """Start tracing.
+
+    Raises:
+      RuntimeError: If trace already in progress.
+    """
+    if self._battor_process is not None:
+      raise RuntimeError('Trace already in progress')
+    self._battery_utils.SetCharging(False)
+    try:
+      self._battor_process = subprocess.Popen(
+          ['battor_agent',
+           '--battor-path='+self._devstring],
+          stdin=subprocess.PIPE,
+          stdout=subprocess.PIPE,
+          shell=False)
+    except OSError:
+      raise RuntimeError('Error executing battor_agent command.'
+                         'Make sure that battor_agent is on PATH.')
+    self._battor_process.stdin.write('StartTracing\n')
+    status = self._battor_process.stdout.readline()
+    if 'Done.' not in status:
+      raise RuntimeError('battor_agent failed to start (%s)' % status)
+    return True
+
+  def _record_trace_result(self):
+    result = self._battor_process.communicate('StopTracing\n')[0]
+    self._battor_process = None
+    self._battery_utils.SetCharging(True)
+    if not result:
+      self._recording_error = 'battor_agent failed to stop - no output'
+    elif result[-6:-1] != 'Done.':
+      self._recording_error = 'battor_agent failed to stop'
+    else:
+      self._recording_error = None
+      self._trace_result = result
+
+  def StopAgentTracing(self):
+    """Collect the result of tracing.
+    """
+    if self._battor_process is not None:
+      self._collection_process = threading.Thread(
+          target = self._record_trace_result)
+      self._collection_process.start()
+      return True
+    else:
+      return False
+
+  def SupportsExplicitClockSync(self):
+    """Returns whether this function supports explicit clock sync."""
+    self._battor_process.stdin.write('SupportsExplicitClockSync\n')
+    result = self._battor_process.stdout.readline()
+    return int(result) != 0
+
+  def RecordClockSyncMarker(self, sync_id, did_record_sync_marker_callback):
+    """Records a clock sync marker.
+
+    Args:
+        sync_id: ID string for clock sync marker.
+    """
+    t1 = trace_time.Now()
+    self._battor_process.stdin.write('RecordClockSyncMarker %s\n' % sync_id)
+    status = self._battor_process.stdout.readline()
+    did_record_sync_marker_callback(t1, sync_id)
+    if 'Done.' not in status:
+      raise RuntimeError('battor_agent failed to record sync marker (%s)'
+                        % status)
+
+  def GetResults(self):
+    """Get the trace data.
+
+    Returns:
+      The trace data.
+    """
+    if self._collection_process is None:
+      raise RuntimeError('Must call StopAgentTracing before GetResults')
+    self._collection_process.join()
+    self._collection_process = None
+    if self._recording_error:
+      raise RuntimeError(self._recording_error)
+    else:
+      return TraceResults('powerTraceAsString', self._trace_result, True)