Move RegisterDomain/UnregisterDomain and its handling to inspector_websocket.

tools/telemetry/telemetry/core/backends/chrome/inspector_websocket.py

Index: tools/telemetry/telemetry/core/backends/chrome/inspector_websocket.py
diff --git a/tools/telemetry/telemetry/core/backends/chrome/inspector_websocket.py b/tools/telemetry/telemetry/core/backends/chrome/inspector_websocket.py
index 899bc9c4723b4e9a909e59551a93275f27b8edb1..e478dc690a7222d871d0627165bdf59450e14f47 100644
--- a/tools/telemetry/telemetry/core/backends/chrome/inspector_websocket.py
+++ b/tools/telemetry/telemetry/core/backends/chrome/inspector_websocket.py
@@ -2,6 +2,7 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
+import collections
 import json
 import logging
 import socket
@@ -10,6 +11,10 @@ import time
 from telemetry.core.backends.chrome import websocket
 
 
+_DomainHandler = collections.namedtuple(
+    'DomainHandler', ['notification_handler', 'will_close_handler'])
+
+
 class DispatchNotificationsUntilDoneTimeoutException(Exception):
   """Exception that can be thrown from DispatchNotificationsUntilDone to
   indicate timeout exception of the function.
@@ -22,15 +27,10 @@ class DispatchNotificationsUntilDoneTimeoutException(Exception):
 
 class InspectorWebsocket(object):
 
-  def __init__(self, notification_handler=None, error_handler=None):
+  def __init__(self, error_handler=None):
     """Create a websocket handler for communicating with Inspectors.
 
     Args:
-      notification_handler: A callback for notifications received as a result of
-        calling DispatchNotifications() or DispatchNotificationsUntilDone().
-        Must accept a single JSON object containing the Inspector's
-        notification. May return True to indicate the dispatching is done for
-        DispatchNotificationsUntilDone.
       error_handler: A callback for errors in communicating with the Inspector.
         Must accept a single numeric parameter indicated the time elapsed before
         the error.
@@ -38,9 +38,45 @@ class InspectorWebsocket(object):
     self._socket = None
     self._cur_socket_timeout = 0
     self._next_request_id = 0
-    self._notification_handler = notification_handler
     self._error_handler = error_handler
     self._all_data_received = False
+    self._domain_handlers = {}
+
+  def RegisterDomain(
+      self, domain_name, notification_handler, will_close_handler=None):
+    """Registers a given domain for handling notification methods.
+
+    When used as handler for DispatchNotificationsUntilDone,
+    notification handler should return a boolean, where True indicates
+    that we should stop listening for more notifications.
+
+    For example, given inspector_backend:
+       def OnConsoleNotification(msg):
+          if msg['method'] == 'Console.messageAdded':
+             print msg['params']['message']
+          return True
+       def OnConsoleClose(self):
+          pass
+       inspector_backend.RegisterDomain(
+           'Console', OnConsoleNotification, OnConsoleClose)
+
+    Args:
+      domain_name: The devtools domain name. E.g., 'Tracing', 'Memory', 'Page'.
+      notification_handler: Handler for devtools notification. Will be
+          called if a devtools notification with matching domain is received
+          (via DispatchNotifications and DispatchNotificationsUntilDone).
+          The handler accepts a single paramater: the JSON object representing
+          the notification.
+      will_close_handler: Handler to be called from Disconnect().
+    """
+    assert domain_name not in self._domain_handlers
+    self._domain_handlers[domain_name] = _DomainHandler(
+        notification_handler, will_close_handler)
+
+  def UnregisterDomain(self, domain_name):
+    """Unregisters a previously registered domain."""
+    assert domain_name in self._domain_handlers
+    self._domain_handlers.pop(domain_name)
 
   def Connect(self, url, timeout=10):
     assert not self._socket
@@ -49,6 +85,14 @@ class InspectorWebsocket(object):
     self._next_request_id = 0
 
   def Disconnect(self):
+    """Disconnects the inspector websocket.
+
+    All existing domain handlers will also be unregistered.
+    """
+    for _, handler in self._domain_handlers.items():
+      if handler.will_close_handler:
+        handler.will_close_handler()
+
     if self._socket:
       self._socket.close()
       self._socket = None
@@ -95,6 +139,10 @@ class InspectorWebsocket(object):
         if self._all_data_received:
           break
       except websocket.WebSocketTimeoutException:
+        # TODO(chrishenry): Since we always call settimeout in
+        # _Receive, we should be able to rip manual logic of tracking
+        # elapsed time and simply throw
+        # DispatchNotificationsUntilDoneTimeoutException from here.
         pass
       elapsed_time = time.time() - timeout_start_time
       if elapsed_time > timeout:
@@ -112,13 +160,24 @@ class InspectorWebsocket(object):
       if self._socket:
         self._all_data_received = False
         data = self._socket.recv()
-        res = json.loads(data)
+        result = json.loads(data)
         if logging.getLogger().isEnabledFor(logging.DEBUG):
-          logging.debug('got [%s]', json.dumps(res, indent=2, sort_keys=True))
-        if 'method' in res and self._notification_handler(res):
+          logging.debug(
+              'got [%s]', json.dumps(result, indent=2, sort_keys=True))
+        if 'method' in result and self._HandleNotification(result):
           self._all_data_received = True
           return None
-        return res
+        return result
     except (socket.error, websocket.WebSocketException):
       elapsed_time = time.time() - start_time
       self._error_handler(elapsed_time)
+
+  def _HandleNotification(self, result):
+    mname = result['method']
+    dot_pos = mname.find('.')
+    domain_name = mname[:dot_pos]
+    if domain_name in self._domain_handlers:
+      return self._domain_handlers[domain_name].notification_handler(result)
+
+    logging.warn('Unhandled inspector message: %s', result)
+    return False