[Telemetry] Add battor tracing agent

telemetry/telemetry/internal/platform/tracing_agent/battor_tracing_agent_unittest.py

Index: telemetry/telemetry/internal/platform/tracing_agent/battor_tracing_agent_unittest.py
diff --git a/telemetry/telemetry/internal/platform/tracing_agent/battor_tracing_agent_unittest.py b/telemetry/telemetry/internal/platform/tracing_agent/battor_tracing_agent_unittest.py
new file mode 100644
index 0000000000000000000000000000000000000000..f91903403b3826bb9ede891e662397ef77d0c065
--- /dev/null
+++ b/telemetry/telemetry/internal/platform/tracing_agent/battor_tracing_agent_unittest.py
@@ -0,0 +1,168 @@
+# Copyright 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 unittest
+
+from battor import battor_error
+from battor import battor_wrapper
+from telemetry.internal.platform.tracing_agent import battor_tracing_agent
+from telemetry.timeline import trace_data
+from telemetry.timeline import tracing_config
+
+_BATTOR_RETURN = ['fake', 'battor', 'data']
+
+
+class FakePlatformBackend(object):
+  def GetOSName(self):
+    return ''
+
+
+class FakeAndroidPlatformBackend(FakePlatformBackend):
+  def __init__(self):
+    super(FakeAndroidPlatformBackend, self).__init__()
+    self.device = 'fake_device'
+
+  def GetOSName(self):
+    return 'android'
+
+
+class FakeDesktopPlatformBackend(FakePlatformBackend):
+  def __init__(self):
+    self.platform = 'win'
+
+  def GetOSName(self):
+    return self.platform
+
+
+class FakeBattOr(object):
+  def __init__(self, test_platform, android_device=None, battor_path=None,
+               battor_map=None):
+    self._is_shell_running = False
+    self._android_device = android_device
+    self._battor_path = battor_path
+    self._battor_map = battor_map
+    self._test_platform = test_platform
+    self._stop_tracing_called = False
+    self._start_shell_called = False
+    self._start_tracing_called = False
+    self._collect_trace_data_called = False
+    self._record_clock_sync_marker_called = False
+
+  def IsShellRunning(self):
+    return self._is_shell_running
+
+  def StartShell(self):
+    self._is_shell_running = True
+    self._start_shell_called = True
+
+  def StartTracing(self):
+    self.StartShell()
+    self._start_tracing_called = True
+
+  def StopTracing(self):
+    self._is_shell_running = False
+    self._stop_tracing_called = True
+
+  def CollectTraceData(self, timeout=None):
+    del timeout # unused
+    self._collect_trace_data_called = True
+    return _BATTOR_RETURN
+
+  def RecordClockSyncMarker(self, _):
+    self._record_clock_sync_marker_called = True
+
+
+class BattOrTracingAgentTest(unittest.TestCase):
+  def setUp(self):
+    self._config = tracing_config.TracingConfig()
+    self._config.enable_battor_trace = True
+
+    # Monkey patch BattorWrapper.
+    self._battor_wrapper = battor_wrapper.BattorWrapper
+    battor_wrapper.BattorWrapper = FakeBattOr
+    battor_wrapper.IsBattOrConnected = lambda x, android_device=None: True
+
+    # Agents and backends.
+    self.android_backend = FakeAndroidPlatformBackend()
+    self.desktop_backend = FakeDesktopPlatformBackend()
+    self.android_agent = (
+        battor_tracing_agent.BattOrTracingAgent(self.android_backend))
+    self.desktop_agent = (
+        battor_tracing_agent.BattOrTracingAgent(self.desktop_backend))
+
+  def tearDown(self):
+    battor_wrapper.BattorWrapper = self._battor_wrapper
+
+  def testInit(self):
+    self.assertTrue(isinstance(self.android_agent._platform_backend,
+                               FakeAndroidPlatformBackend))
+    self.assertTrue(isinstance(self.desktop_agent._platform_backend,
+                               FakeDesktopPlatformBackend))
+
+  def testIsSupportedAndroid(self):
+    self.assertTrue(battor_tracing_agent.BattOrTracingAgent.IsSupported(
+        self.android_backend))
+    battor_wrapper.IsBattOrConnected = lambda x, android_device=None: False
+    self.assertFalse(battor_tracing_agent.BattOrTracingAgent.IsSupported(
+        self.android_backend))
+
+  def testIsSupportedNonAndroid(self):
+    self.desktop_backend.platform = 'mac'
+    battor_wrapper.IsBattOrConnected = lambda *unused: True
+    self.assertTrue(battor_tracing_agent.BattOrTracingAgent.IsSupported(
+        self.desktop_backend))
+    battor_wrapper.IsBattOrConnected = lambda *unused: False
+    self.assertFalse(battor_tracing_agent.BattOrTracingAgent.IsSupported(
+        self.desktop_backend))
+
+  def testStartAgentTracingPass(self):
+    self.assertTrue(self.android_agent.StartAgentTracing(self._config, 0))
+    self.assertTrue(self.android_agent._battor._is_shell_running)
+    self.assertTrue(self.android_agent._battor._start_shell_called)
+    self.assertTrue(self.android_agent._battor._start_tracing_called)
+    self.assertFalse(self.android_agent._battor._stop_tracing_called)
+    self.assertFalse(
+        self.android_agent._battor._record_clock_sync_marker_called)
+
+  def testStartAgentTracingConfigSetToFalse(self):
+    self._config.enable_battor_trace = False
+    self.assertFalse(self.android_agent.StartAgentTracing(self._config, 0))
+    self.assertFalse(self.android_agent._battor._is_shell_running)
+    self.assertFalse(self.android_agent._battor._start_shell_called)
+    self.assertFalse(self.android_agent._battor._start_tracing_called)
+    self.assertFalse(self.android_agent._battor._stop_tracing_called)
+    self.assertFalse(
+        self.android_agent._battor._record_clock_sync_marker_called)
+
+  def testStartAgentTracingFail(self):
+    def throw_battor_error():
+      raise battor_error.BattorError('Forced Exception')
+    self.android_agent._battor.StartTracing = throw_battor_error
+    self.assertFalse(self.android_agent.StartAgentTracing(self._config, 0))
+
+  def testStopAgentTracing(self):
+    self.android_agent.StopAgentTracing()
+    self.assertTrue(self.android_agent._battor._stop_tracing_called)
+
+  def testCollectAgentTraceData(self):
+    builder = trace_data.TraceDataBuilder()
+    self.android_agent.CollectAgentTraceData(builder)
+    self.assertTrue(self.android_agent._battor._collect_trace_data_called)
+    builder = builder.AsData()
+    self.assertTrue(builder.HasEventsFor(trace_data.BATTOR_TRACE_PART))
+    data_from_builder = builder.GetEventsFor(trace_data.BATTOR_TRACE_PART)
+    self.assertListEqual(['\n'.join(_BATTOR_RETURN)], data_from_builder)
+
+  def testRecordClockSyncMarker(self):
+    def callback_with_exception(a, b):
+      del a # unused
+      del b # unused
+      raise Exception
+    def callback_without_exception(a, b):
+      del a # unused
+      del b # unused
+
+    self.android_agent.RecordClockSyncMarker('123', callback_without_exception)
+    with self.assertRaises(Exception):
+      self.android_agent.RecordClockSyncMarker('abc', callback_with_exception)