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])