Fixing the api_streams sensor (#22200)

* Fire events with websocket messages.

* Added tests to validate

* Fixed api_streams sensor to use new sensor

* Delete from coverageac as now works.

* Removed websocket request event.

* Use dispatcher instead of events.

* Moved sensor to under websocket_api

* Changes as per code review

* Fixed tests.

* Modified test

* Patch

.coveragerc

@@ -415,7 +415,6 @@ omit =
     homeassistant/components/lifx_cloud/scene.py
     homeassistant/components/scsgate/*
     homeassistant/components/sense/*
-    homeassistant/components/api_streams/sensor.py
     homeassistant/components/aftership/sensor.py
     homeassistant/components/airvisual/sensor.py
     homeassistant/components/alpha_vantage/sensor.py

homeassistant/components/websocket_api/const.py

@@ -20,3 +20,7 @@
 # Originally, this was just asyncio.CancelledError, but issue #9546 showed
 # that futures.CancelledErrors can also occur in some situations.
 CANCELLATION_ERRORS = (asyncio.CancelledError, futures.CancelledError)
+
+# Event types
+SIGNAL_WEBSOCKET_CONNECTED = 'websocket_connected'
+SIGNAL_WEBSOCKET_DISCONNECTED = 'websocket_disconnected'

homeassistant/components/websocket_api/http.py

@@ -13,7 +13,9 @@
 from homeassistant.components.http import HomeAssistantView
 from homeassistant.helpers.json import JSONEncoder
 
-from .const import MAX_PENDING_MSG, CANCELLATION_ERRORS, URL, ERR_UNKNOWN_ERROR
+from .const import (
+    MAX_PENDING_MSG, CANCELLATION_ERRORS, URL, ERR_UNKNOWN_ERROR,
+    SIGNAL_WEBSOCKET_CONNECTED, SIGNAL_WEBSOCKET_DISCONNECTED)
 from .auth import AuthPhase, auth_required_message
 from .error import Disconnect
 from .messages import error_message
@@ -142,6 +144,8 @@ def handle_hass_stop(event):
 
             self._logger.debug("Received %s", msg)
             connection = await auth.async_handle(msg)
+            self.hass.helpers.dispatcher.async_dispatcher_send(
+                SIGNAL_WEBSOCKET_CONNECTED)
 
             # Command phase
             while not wsock.closed:
@@ -192,4 +196,7 @@ def handle_hass_stop(event):
             else:
                 self._logger.warning("Disconnected: %s", disconnect_warn)
 
+            self.hass.helpers.dispatcher.async_dispatcher_send(
+                SIGNAL_WEBSOCKET_DISCONNECTED)
+
         return wsock

homeassistant/components/websocket_api/sensor.py

@@ -0,0 +1,53 @@
+"""Entity to track connections to websocket API."""
+
+from homeassistant.core import callback
+from homeassistant.helpers.entity import Entity
+
+from .const import SIGNAL_WEBSOCKET_CONNECTED, SIGNAL_WEBSOCKET_DISCONNECTED
+
+
+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])
+
+
+class APICount(Entity):
+    """Entity to represent how many people are connected to the stream API."""
+
+    def __init__(self):
+        """Initialize the API count."""
+        self.count = 0
+
+    @property
+    def name(self):
+        """Return name of entity."""
+        return "Connected clients"
+
+    @property
+    def state(self):
+        """Return current API count."""
+        return self.count
+
+    @property
+    def unit_of_measurement(self):
+        """Return the unit of measurement."""
+        return "clients"
+
+    @callback
+    def _increment(self):
+        self.count += 1
+        self.async_schedule_update_ha_state()
+
+    @callback
+    def _decrement(self):
+        self.count -= 1
+        self.async_schedule_update_ha_state()

tests/components/websocket_api/test_auth.py

@@ -1,7 +1,8 @@
 """Test auth of websocket API."""
 from unittest.mock import patch
 
-from homeassistant.components.websocket_api.const import URL
+from homeassistant.components.websocket_api.const import (
+    URL, SIGNAL_WEBSOCKET_CONNECTED, SIGNAL_WEBSOCKET_DISCONNECTED)
 from homeassistant.components.websocket_api.auth import (
     TYPE_AUTH, TYPE_AUTH_INVALID, TYPE_AUTH_OK, TYPE_AUTH_REQUIRED)
 
@@ -24,6 +25,28 @@ async def test_auth_via_msg(no_auth_websocket_client, legacy_auth):
     assert msg['type'] == TYPE_AUTH_OK
 
 
+async def test_auth_events(hass, no_auth_websocket_client, legacy_auth):
+    """Test authenticating."""
+    connected_evt = []
+    hass.helpers.dispatcher.async_dispatcher_connect(
+        SIGNAL_WEBSOCKET_CONNECTED,
+        lambda: connected_evt.append(1))
+    disconnected_evt = []
+    hass.helpers.dispatcher.async_dispatcher_connect(
+        SIGNAL_WEBSOCKET_DISCONNECTED,
+        lambda: disconnected_evt.append(1))
+
+    await test_auth_via_msg(no_auth_websocket_client, legacy_auth)
+
+    assert len(connected_evt) == 1
+    assert not disconnected_evt
+
+    await no_auth_websocket_client.close()
+    await hass.async_block_till_done()
+
+    assert len(disconnected_evt) == 1
+
+
 async def test_auth_via_msg_incorrect_pass(no_auth_websocket_client):
     """Test authenticating."""
     with patch('homeassistant.components.websocket_api.auth.'
@@ -41,6 +64,29 @@ async def test_auth_via_msg_incorrect_pass(no_auth_websocket_client):
     assert msg['message'] == 'Invalid access token or password'
 
 
+async def test_auth_events_incorrect_pass(hass, no_auth_websocket_client):
+    """Test authenticating."""
+    connected_evt = []
+    hass.helpers.dispatcher.async_dispatcher_connect(
+        SIGNAL_WEBSOCKET_CONNECTED,
+        lambda: connected_evt.append(1))
+    disconnected_evt = []
+    hass.helpers.dispatcher.async_dispatcher_connect(
+        SIGNAL_WEBSOCKET_DISCONNECTED,
+        lambda: disconnected_evt.append(1))
+
+    await test_auth_via_msg_incorrect_pass(no_auth_websocket_client)
+
+    assert not connected_evt
+    assert not disconnected_evt
+
+    await no_auth_websocket_client.close()
+    await hass.async_block_till_done()
+
+    assert not connected_evt
+    assert not disconnected_evt
+
+
 async def test_pre_auth_only_auth_allowed(no_auth_websocket_client):
     """Verify that before authentication, only auth messages are allowed."""
     await no_auth_websocket_client.send_json({

tests/components/websocket_api/test_sensor.py

@@ -0,0 +1,30 @@
+"""Test cases for the API stream sensor."""
+
+from homeassistant.bootstrap import async_setup_component
+
+from tests.common import assert_setup_component
+from .test_auth import test_auth_via_msg
+
+
+async def test_websocket_api(hass, no_auth_websocket_client, legacy_auth):
+    """Test API streams."""
+    with assert_setup_component(1):
+        await async_setup_component(hass, 'sensor', {
+            'sensor': {
+                'platform': 'websocket_api',
+            }
+        })
+
+    state = hass.states.get('sensor.connected_clients')
+    assert state.state == '0'
+
+    await test_auth_via_msg(no_auth_websocket_client, legacy_auth)
+
+    state = hass.states.get('sensor.connected_clients')
+    assert state.state == '1'
+
+    await no_auth_websocket_client.close()
+    await hass.async_block_till_done()
+
+    state = hass.states.get('sensor.connected_clients')
+    assert state.state == '0'