[Systrace] Fix systrace clock syncing issue with BattOr.

systrace/systrace/tracing_agents/battor_trace_agent.py

Index: systrace/systrace/tracing_agents/battor_trace_agent.py
diff --git a/systrace/systrace/tracing_agents/battor_trace_agent.py b/systrace/systrace/tracing_agents/battor_trace_agent.py
index 80cd274d25d99000207ba39c920a70b96e67bc7e..6e7d9743373c38d737e86939f88b9c0423d37e6d 100644
--- a/systrace/systrace/tracing_agents/battor_trace_agent.py
+++ b/systrace/systrace/tracing_agents/battor_trace_agent.py
@@ -2,7 +2,6 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
-from os import path
 import atexit
 import logging
 import optparse
@@ -12,9 +11,11 @@ from battor import battor_wrapper
 from devil.android import battery_utils
 from devil.android import device_utils
 from devil.utils import battor_device_mapping
+from devil.utils import find_usb_devices
 from py_trace_event import trace_time
 from systrace import trace_result
 from systrace import tracing_agents
+from telemetry.timeline import trace_data
 
 
 def try_create_agent(config):
@@ -26,45 +27,37 @@ def try_create_agent(config):
 
 
 class BattOrConfig(tracing_agents.TracingConfig):
-  def __init__(self, battor_categories, hub_types, serial_map, battor_path,
-               update_map, battor, target, from_file):
+  def __init__(self, battor_categories, serial_map, battor_path,
+                battor, target, from_file, device_serial_number):
     tracing_agents.TracingConfig.__init__(self)
     self.battor_categories = battor_categories
-    self.hub_types = hub_types
     self.serial_map = serial_map
     self.battor_path = battor_path
-    self.update_map = update_map
     self.battor = battor
     self.target = target
     self.from_file = from_file
-
+    self.device_serial_number = device_serial_number
 
 def add_options(parser):
   options = optparse.OptionGroup(parser, 'BattOr trace options')
   options.add_option('--battor-categories', dest='battor_categories',
                      help='Select battor categories with a comma-delimited '
                      'list, e.g. --battor-categories=cat1,cat2,cat3')
-  options.add_option('--hubs', dest='hub_types', default='plugable_7port',
-                    help='List of hub types to check for for BattOr mapping. '
-                    'Used when updating mapping file.')
   options.add_option('--serial-map', dest='serial_map',
                     default='serial_map.json',
                     help='File containing pregenerated map of phone serial '
                     'numbers to BattOr serial numbers.')
-  options.add_option('--battor_path', dest='battor_path', default=None,
+  options.add_option('--battor-path', dest='battor_path', default=None,
                     type='string', help='specify a BattOr path to use')
-  options.add_option('--update-map', dest='update_map', default=False,
-                    action='store_true',
-                    help='force update of phone-to-BattOr map')
   options.add_option('--battor', dest='battor', default=False,
                     action='store_true', help='Use the BattOr tracing agent.')
   return options
 
 def get_config(options):
-  return BattOrConfig(options.battor_categories, options.hub_types,
+  return BattOrConfig(options.battor_categories,
                       options.serial_map, options.battor_path,
-                      options.update_map, options.battor, options.target,
-                      options.from_file)
+                      options.battor, options.target,
+                      options.from_file, options.device_serial_number)
 
 def _reenable_charging_if_needed(battery):
   if not battery.GetCharging():
@@ -92,13 +85,17 @@ class BattOrTraceAgent(tracing_agents.TracingAgent):
     Raises:
         RuntimeError: If trace already in progress.
     """
-    if config.update_map or not path.isfile(config.serial_map):
-      battor_device_mapping.GenerateSerialMapFile(config.serial_map,
-                                                  config.hub_types)
+    device_tree = find_usb_devices.GetBusNumberToDeviceTreeMap()
+    battors = battor_device_mapping.GetBattOrList(device_tree)
+    battor = config.battor_path
+    if not config.battor_path:
+      assert len(battors) == 1
+      battor = battors[0]
+
     self._battor_wrapper = battor_wrapper.BattOrWrapper(
         target_platform=config.target,
         android_device=config.device_serial_number,
-        battor_path=config.battor_path,
+        battor_path=battor,
         battor_map_file=config.serial_map)
 
     dev_utils = device_utils.DeviceUtils(config.device_serial_number)
@@ -159,3 +156,7 @@ class BattOrTraceAgent(tracing_agents.TracingAgent):
     """
     return trace_result.TraceResult(
         'powerTraceAsString', self._battor_wrapper.CollectTraceData())
+
+  def CollectAgentTraceData(self, trace_data_builder):
+    data = self._battor_wrapper.CollectTraceData()
+    trace_data_builder.AddTraceFor(trace_data.BATTOR_TRACE_PART, data)