[Systrace][BattOr] Simplify BattOr options in systrace.

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..a48ea87ee4cc2a04da020c2258995984b9013e1e 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,6 +11,7 @@ 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
@@ -26,17 +26,16 @@ 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):
@@ -44,27 +43,21 @@ def add_options(parser):
   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,
-                      options.serial_map, options.battor_path,
-                      options.update_map, options.battor, options.target,
-                      options.from_file)
+  return BattOrConfig(
+      options.battor_categories, options.serial_map, options.battor_path,
+      options.battor, options.target, options.from_file,
+      options.device_serial_number)
 
 def _reenable_charging_if_needed(battery):
   if not battery.GetCharging():
@@ -82,6 +75,17 @@ class BattOrTraceAgent(tracing_agents.TracingAgent):
     self._battor_wrapper = None
     self._battery_utils = None
 
+  @staticmethod
+  def _FindBattOrPath(config):
+    device_tree = find_usb_devices.GetBusNumberToDeviceTreeMap()
+    battors = battor_device_mapping.GetBattOrList(device_tree)
+    battor_path = config.battor_path
+    if not config.battor_path and not config.serial_map:
+      assert len(battors) == 1, ('Must specify BattOr path if there is not '
+                                 'exactly one')
+      battor_path = battors[0]
+    return battor_path
+
   @py_utils.Timeout(tracing_agents.START_STOP_TIMEOUT)
   def StartAgentTracing(self, config, timeout=None):
     """Starts tracing.
@@ -91,14 +95,14 @@ class BattOrTraceAgent(tracing_agents.TracingAgent):
 
     Raises:
         RuntimeError: If trace already in progress.
+        AssertionError: If There is no BattOr path given and more
+            than one BattOr is attached.
     """
-    if config.update_map or not path.isfile(config.serial_map):
-      battor_device_mapping.GenerateSerialMapFile(config.serial_map,
-                                                  config.hub_types)
+    battor_path = self._FindBattOrPath(config)
     self._battor_wrapper = battor_wrapper.BattOrWrapper(
         target_platform=config.target,
         android_device=config.device_serial_number,
-        battor_path=config.battor_path,
+        battor_path=battor_path,
         battor_map_file=config.serial_map)
 
     dev_utils = device_utils.DeviceUtils(config.device_serial_number)