From 2527731865b2be5944f8fed39872a5a4283a4f21 Mon Sep 17 00:00:00 2001 From: Penny Wood Date: Sun, 14 Apr 2019 01:48:40 +0800 Subject: [PATCH] Fix websocket connection sensor (#22923) * Fix for #22890 * Singleton count --- .../components/websocket_api/const.py | 3 ++ .../components/websocket_api/http.py | 6 +++- .../components/websocket_api/sensor.py | 29 +++++++++---------- 3 files changed, 22 insertions(+), 16 deletions(-) diff --git a/homeassistant/components/websocket_api/const.py b/homeassistant/components/websocket_api/const.py index 4c3e0d564dc..53ca680c4c9 100644 --- a/homeassistant/components/websocket_api/const.py +++ b/homeassistant/components/websocket_api/const.py @@ -24,3 +24,6 @@ CANCELLATION_ERRORS = (asyncio.CancelledError, futures.CancelledError) # Event types SIGNAL_WEBSOCKET_CONNECTED = 'websocket_connected' SIGNAL_WEBSOCKET_DISCONNECTED = 'websocket_disconnected' + +# Data used to store the current connection list +DATA_CONNECTIONS = DOMAIN + '.connections' diff --git a/homeassistant/components/websocket_api/http.py b/homeassistant/components/websocket_api/http.py index b51e5d5699d..0fc446390b7 100644 --- a/homeassistant/components/websocket_api/http.py +++ b/homeassistant/components/websocket_api/http.py @@ -15,7 +15,8 @@ from homeassistant.helpers.json import JSONEncoder 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 @@ -145,6 +146,8 @@ class WebSocketHandler: 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) @@ -197,6 +200,7 @@ class WebSocketHandler: else: self._logger.warning("Disconnected: %s", disconnect_warn) + self.hass.data[DATA_CONNECTIONS] -= 1 self.hass.helpers.dispatcher.async_dispatcher_send( SIGNAL_WEBSOCKET_DISCONNECTED) diff --git a/homeassistant/components/websocket_api/sensor.py b/homeassistant/components/websocket_api/sensor.py index fd9108c0513..b43e356b9ce 100644 --- a/homeassistant/components/websocket_api/sensor.py +++ b/homeassistant/components/websocket_api/sensor.py @@ -3,7 +3,9 @@ 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( @@ -11,12 +13,6 @@ async def async_setup_platform( """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 @@ class APICount(Entity): 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()