Telemetry: Refactor inspector_websocket to create cleaner interface.

tools/telemetry/telemetry/core/backends/chrome_inspector/tracing_backend_unittest.py

Index: tools/telemetry/telemetry/core/backends/chrome_inspector/tracing_backend_unittest.py
diff --git a/tools/telemetry/telemetry/core/backends/chrome_inspector/tracing_backend_unittest.py b/tools/telemetry/telemetry/core/backends/chrome_inspector/tracing_backend_unittest.py
index 48d8650fe8fc3616c104fdaf47826866a42e5ff0..115d7dec2d9c8486d272e68343e82594cabde998 100644
--- a/tools/telemetry/telemetry/core/backends/chrome_inspector/tracing_backend_unittest.py
+++ b/tools/telemetry/telemetry/core/backends/chrome_inspector/tracing_backend_unittest.py
@@ -8,12 +8,62 @@ import unittest
 
 from telemetry import decorators
 from telemetry.core import util
+from telemetry.core.backends.chrome_inspector import inspector_websocket
+from telemetry.core.backends.chrome_inspector import tracing_backend
+from telemetry.core.backends.chrome_inspector import websocket
 from telemetry.core.platform import tracing_category_filter
 from telemetry.core.platform import tracing_options
 from telemetry.timeline import model as model_module
 from telemetry.timeline import trace_data as trace_data_module
+from telemetry.unittest_util import simple_mock
 from telemetry.unittest_util import tab_test_case
 
+util.AddDirToPythonPath(util.GetTelemetryDir(), 'third_party', 'mock')
+import mock
+
+
+class FakeInspectorWebsocket(object):
+  """A fake InspectorWebsocket.
+
+  A fake that allows tests to send pregenerated data. Normal
+  InspectorWebsockets allow for any number of domain handlers. This fake only
+  allows up to 1 domain handler, and assumes that the domain of the response
+  always matches that of the handler.
+  """
+  def __init__(self, mock_timer):
+    self._mock_timer = mock_timer
+    self._responses = []
+    self._handler = None
+
+  def RegisterDomain(self, _, handler):
+    self._handler = handler
+
+  def AddResponse(self, method, value, time):
+    if self._responses:
+      assert self._responses[-1][1] < time, (
+          'Current response is scheduled earlier than previous response.')
+    params = {'value': value}
+    response = {'method': method, 'params': params}
+    self._responses.append((response, time))
+
+  def Connect(self, _):
+    pass
+
+  def DispatchNotifications(self, timeout):
+    current_time = self._mock_timer.time()
+    if not self._responses:
+      self._mock_timer.SetTime(current_time + timeout + 1)
+      raise websocket.WebSocketTimeoutException()
+
+    response, time = self._responses[0]
+    if time - current_time > timeout:
+      self._mock_timer.SetTime(current_time + timeout + 1)
+      raise websocket.WebSocketTimeoutException()
+
+    self._responses.pop(0)
+    self._mock_timer.SetTime(time + 1)
+    self._handler(response)
+
 
 class TracingBackendTest(tab_test_case.TabTestCase):
 
@@ -53,3 +103,44 @@ class TracingBackendTest(tab_test_case.TabTestCase):
     self.assertFalse(self._tracing_controller.is_tracing_running)
     # Calling stop again will raise exception
     self.assertRaises(Exception, self._tracing_controller.Stop)
+
+
+class TracingBackendUnitTest(unittest.TestCase):
+  def setUp(self):
+    self._mock_timer = simple_mock.MockTimer(tracing_backend)
+
+  def tearDown(self):
+    self._mock_timer.Restore()
+
+  def testCollectTracingDataTimeout(self):
+    inspector = FakeInspectorWebsocket(self._mock_timer)
+    inspector.AddResponse('Tracing.dataCollected', 'asdf1', 9)
+    inspector.AddResponse('Tracing.dataCollected', 'asdf2', 19)
+    inspector.AddResponse('Tracing.tracingComplete', 'asdf3', 35)
+
+    with mock.patch('telemetry.core.backends.chrome_inspector.'
+                    'inspector_websocket.InspectorWebsocket') as mock_class:
+      mock_class.return_value = inspector
+      backend = tracing_backend.TracingBackend(devtools_port=65000)
+
+    # The third response is 16 seconds after the second response, so we expect
+    # a TracingTimeoutException.
+    with self.assertRaises(tracing_backend.TracingTimeoutException):
+      backend._CollectTracingData(10)
+    self.assertEqual(2, len(backend._trace_events))
+    self.assertFalse(backend._has_received_all_tracing_data)
+
+  def testCollectTracingDataNoTimeout(self):
+    inspector = FakeInspectorWebsocket(self._mock_timer)
+    inspector.AddResponse('Tracing.dataCollected', 'asdf1', 9)
+    inspector.AddResponse('Tracing.dataCollected', 'asdf2', 14)
+    inspector.AddResponse('Tracing.tracingComplete', 'asdf3', 19)
+
+    with mock.patch('telemetry.core.backends.chrome_inspector.'
+                    'inspector_websocket.InspectorWebsocket') as mock_class:
+      mock_class.return_value = inspector
+      backend = tracing_backend.TracingBackend(devtools_port=65000)
+
+    backend._CollectTracingData(10)
+    self.assertEqual(2, len(backend._trace_events))
+    self.assertTrue(backend._has_received_all_tracing_data)