Have api_streams sensor also monitor websocket connections (#4668)

homeassistant/components/sensor/api_streams.py

@@ -2,9 +2,15 @@
 import asyncio
 import logging
 
+from homeassistant.const import EVENT_HOMEASSISTANT_STOP
+from homeassistant.core import callback
 from homeassistant.helpers.entity import Entity
 
 
+NAME_WS = 'homeassistant.components.websocket_api'
+NAME_STREAM = 'homeassistant.components.api'
+
+
 class StreamHandler(logging.Handler):
     """Check log messages for stream connect/disconnect."""
 
@@ -16,13 +22,24 @@ class StreamHandler(logging.Handler):
 
     def handle(self, record):
         """Handle a log message."""
-        if not record.msg.startswith('STREAM'):
-            return
+        if record.name == NAME_STREAM:
+            if not record.msg.startswith('STREAM'):
+                return
 
-        if record.msg.endswith('ATTACHED'):
-            self.entity.count += 1
-        elif record.msg.endswith('RESPONSE CLOSED'):
-            self.entity.count -= 1
+            if record.msg.endswith('ATTACHED'):
+                self.entity.count += 1
+            elif record.msg.endswith('RESPONSE CLOSED'):
+                self.entity.count -= 1
+
+        else:
+            if not record.msg.startswith('WS'):
+                return
+            elif len(record.args) < 2:
+                return
+            elif record.args[1] == 'Connected':
+                self.entity.count += 1
+            elif record.args[1] == 'Closed connection':
+                self.entity.count -= 1
 
         self.entity.schedule_update_ha_state()
 
@@ -31,9 +48,18 @@ class StreamHandler(logging.Handler):
 def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
     """Set up the logger for filters."""
     entity = APICount()
+    handler = StreamHandler(entity)
 
-    logging.getLogger('homeassistant.components.api').addHandler(
-        StreamHandler(entity))
+    logging.getLogger(NAME_STREAM).addHandler(handler)
+    logging.getLogger(NAME_WS).addHandler(handler)
+
+    @callback
+    def remove_logger(event):
+        """Remove our handlers."""
+        logging.getLogger(NAME_STREAM).removeHandler(handler)
+        logging.getLogger(NAME_WS).removeHandler(handler)
+
+    hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, remove_logger)
 
     yield from async_add_devices([entity])
 

tests/components/sensor/test_api_streams.py

@@ -37,3 +37,37 @@ def test_api_streams(hass):
 
     state = hass.states.get('sensor.connected_clients')
     assert state.state == '1'
+
+
+@asyncio.coroutine
+def test_websocket_api(hass):
+    """Test API streams."""
+    log = logging.getLogger('homeassistant.components.websocket_api')
+
+    with assert_setup_component(1):
+        yield from async_setup_component(hass, 'sensor', {
+            'sensor': {
+                'platform': 'api_streams',
+            }
+        })
+
+    state = hass.states.get('sensor.connected_clients')
+    assert state.state == '0'
+
+    log.debug('WS %s: %s', id(log), 'Connected')
+    yield from hass.async_block_till_done()
+
+    state = hass.states.get('sensor.connected_clients')
+    assert state.state == '1'
+
+    log.debug('WS %s: %s', id(log), 'Connected')
+    yield from hass.async_block_till_done()
+
+    state = hass.states.get('sensor.connected_clients')
+    assert state.state == '2'
+
+    log.debug('WS %s: %s', id(log), 'Closed connection')
+    yield from hass.async_block_till_done()
+
+    state = hass.states.get('sensor.connected_clients')
+    assert state.state == '1'