Fix websocket connection sensor

homeassistant/components/websocket_api/const.py

@@ -24,3 +24,6 @@
 # Event types
 SIGNAL_WEBSOCKET_CONNECTED = 'websocket_connected'
 SIGNAL_WEBSOCKET_DISCONNECTED = 'websocket_disconnected'
+
+# Data used to store the current connection list
+DATA_CONNECTIONS = DOMAIN + '.connections'

homeassistant/components/websocket_api/http.py

@@ -15,7 +15,8 @@
 
 from .const import (
     MAX_PENDING_MSG, CANCELLATION_ERRORS, URL, ERR_UNKNOWN_ERROR,
-    SIGNAL_WEBSOCKET_CONNECTED, SIGNAL_WEBSOCKET_DISCONNECTED)
+    SIGNAL_WEBSOCKET_CONNECTED, SIGNAL_WEBSOCKET_DISCONNECTED,
+    DATA_CONNECTIONS)
 from .auth import AuthPhase, auth_required_message
 from .error import Disconnect
 from .messages import error_message
@@ -144,6 +145,8 @@ def handle_hass_stop(event):
 
             self._logger.debug("Received %s", msg)
             connection = await auth.async_handle(msg)
+            self.hass.data[DATA_CONNECTIONS] = \
+                self.hass.data.get(DATA_CONNECTIONS, 0) + 1
             self.hass.helpers.dispatcher.async_dispatcher_send(
                 SIGNAL_WEBSOCKET_CONNECTED)
 
@@ -196,6 +199,7 @@ def handle_hass_stop(event):
             else:
                 self._logger.warning("Disconnected: %s", disconnect_warn)
 
+            self.hass.data[DATA_CONNECTIONS] -= 1
             self.hass.helpers.dispatcher.async_dispatcher_send(
                 SIGNAL_WEBSOCKET_DISCONNECTED)
 

homeassistant/components/websocket_api/sensor.py

@@ -3,20 +3,16 @@
 from homeassistant.core import callback
 from homeassistant.helpers.entity import Entity
 
-from .const import SIGNAL_WEBSOCKET_CONNECTED, SIGNAL_WEBSOCKET_DISCONNECTED
+from .const import (
+    SIGNAL_WEBSOCKET_CONNECTED, SIGNAL_WEBSOCKET_DISCONNECTED,
+    DATA_CONNECTIONS)
 
 
 async def async_setup_platform(
         hass, config, async_add_entities, discovery_info=None):
     """Set up the API streams platform."""
     entity = APICount()
 
-    # pylint: disable=protected-access
-    hass.helpers.dispatcher.async_dispatcher_connect(
-        SIGNAL_WEBSOCKET_CONNECTED, entity._increment)
-    hass.helpers.dispatcher.async_dispatcher_connect(
-        SIGNAL_WEBSOCKET_DISCONNECTED, entity._decrement)
-
     async_add_entities([entity])
 
 
@@ -25,7 +21,15 @@ class APICount(Entity):
 
     def __init__(self):
         """Initialize the API count."""
-        self.count = 0
+        self.count = None
+
+    async def async_added_to_hass(self):
+        """Added to hass."""
+        self.hass.helpers.dispatcher.async_dispatcher_connect(
+            SIGNAL_WEBSOCKET_CONNECTED, self._update_count)
+        self.hass.helpers.dispatcher.async_dispatcher_connect(
+            SIGNAL_WEBSOCKET_DISCONNECTED, self._update_count)
+        self._update_count()
 
     @property
     def name(self):
@@ -43,11 +47,6 @@ def unit_of_measurement(self):
         return "clients"
 
     @callback
-    def _increment(self):
-        self.count += 1
-        self.async_schedule_update_ha_state()
-
-    @callback
-    def _decrement(self):
-        self.count -= 1
+    def _update_count(self):
+        self.count = self.hass.data.get(DATA_CONNECTIONS, 0)
         self.async_schedule_update_ha_state()