1 files changed, 78 insertions, 0 deletions
diff --git a/tools/android/adb_profile_chrome/profiler_unittest.py b/tools/android/adb_profile_chrome/profiler_unittest.py
new file mode 100644
index 0000000..819c440
--- /dev/null
+++ b/tools/android/adb_profile_chrome/profiler_unittest.py
@@ -0,0 +1,78 @@
+# Copyright 2014 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 os
+import tempfile
+import unittest
+import zipfile
+
+from adb_profile_chrome import profiler
+from adb_profile_chrome import ui
+
+
+class FakeController(object):
+  def __init__(self, contents='fake-contents'):
+    self.contents = contents
+    self.interval = None
+    self.stopped = False
+    self.filename = None
+
+  def StartTracing(self, interval):
+    self.interval = interval
+
+  def StopTracing(self):
+    self.stopped = True
+
+  def PullTrace(self):
+    with tempfile.NamedTemporaryFile(delete=False) as f:
+      self.filename = f.name
+      f.write(self.contents)
+      return f.name
+
+  def __repr__(self):
+    return 'faketrace'
+
+
+class ProfilerTest(unittest.TestCase):
+  def setUp(self):
+    ui.EnableTestMode()
+
+  def testCaptureBasicProfile(self):
+    controller = FakeController()
+    interval = 1.5
+    result = profiler.CaptureProfile([controller], interval)
+
+    try:
+      self.assertEquals(controller.interval, interval)
+      self.assertTrue(controller.stopped)
+      self.assertTrue(os.path.exists(result))
+      self.assertFalse(os.path.exists(controller.filename))
+      self.assertTrue(result.endswith('.html'))
+    finally:
+      os.remove(result)
+
+  def testCaptureJsonProfile(self):
+    controller = FakeController()
+    result = profiler.CaptureProfile([controller], 1, write_json=True)
+
+    try:
+      self.assertFalse(result.endswith('.html'))
+      with open(result) as f:
+        self.assertEquals(f.read(), controller.contents)
+    finally:
+      os.remove(result)
+
+  def testCaptureMultipleProfiles(self):
+    controllers = [FakeController('c1'), FakeController('c2')]
+    result = profiler.CaptureProfile(controllers, 1, write_json=True)
+
+    try:
+      self.assertTrue(result.endswith('.zip'))
+      self.assertTrue(zipfile.is_zipfile(result))
+      with zipfile.ZipFile(result) as f:
+        self.assertEquals(
+            f.namelist(),
+            [controllers[0].filename[1:], controllers[1].filename[1:]])
+    finally:
+      os.remove(result)