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/api_streams/sensor.py

@@ -1,90 +0,0 @@
-"""
-Entity to track connections to stream API.
-
-For more details about this platform, please refer to the documentation at
-https://home-assistant.io/components/sensor.api_streams/
-"""
-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."""
-
-    def __init__(self, entity):
-        """Initialize handler."""
-        super().__init__()
-        self.entity = entity
-        self.count = 0
-
-    def handle(self, record):
-        """Handle a log message."""
-        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
-
-        else:
-            if not record.msg.startswith('WS'):
-                return
-            if len(record.args) < 2:
-                return
-            if record.args[1] == 'Connected':
-                self.entity.count += 1
-            elif record.args[1] == 'Closed connection':
-                self.entity.count -= 1
-
-        self.entity.schedule_update_ha_state()
-
-
-async def async_setup_platform(
-        hass, config, async_add_entities, discovery_info=None):
-    """Set up the API streams platform."""
-    entity = APICount()
-    handler = 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)
-
-    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"

homeassistant/components/websocket_api/const.py

@@ -20,3 +20,7 @@ TYPE_RESULT = 'result'
 # 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.core import callback
 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 @@ class WebSocketHandler:
 
             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 @@ class WebSocketHandler:
             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/api_streams/test_sensor.py

@@ -1,75 +0,0 @@
-"""Test cases for the API stream sensor."""
-import asyncio
-import logging
-import pytest
-
-from homeassistant.bootstrap import async_setup_component
-from tests.common import assert_setup_component
-
-
-@pytest.mark.skip(reason="test fails randomly due to race condition.")
-async def test_api_streams(hass):
-    """Test API streams."""
-    log = logging.getLogger('homeassistant.components.api')
-
-    with assert_setup_component(1):
-        await async_setup_component(hass, 'sensor', {
-            'sensor': {
-                'platform': 'api_streams',
-            }
-        })
-
-    state = hass.states.get('sensor.connected_clients')
-    assert state.state == '0'
-
-    log.debug('STREAM 1 ATTACHED')
-    await asyncio.sleep(0.1)
-
-    state = hass.states.get('sensor.connected_clients')
-    assert state.state == '1'
-
-    log.debug('STREAM 1 ATTACHED')
-    await asyncio.sleep(0.1)
-
-    state = hass.states.get('sensor.connected_clients')
-    assert state.state == '2'
-
-    log.debug('STREAM 1 RESPONSE CLOSED')
-    await asyncio.sleep(0.1)
-
-    state = hass.states.get('sensor.connected_clients')
-    assert state.state == '1'
-
-
-@pytest.mark.skip(reason="test fails randomly due to race condition.")
-async def test_websocket_api(hass):
-    """Test API streams."""
-    log = logging.getLogger('homeassistant.components.websocket_api')
-
-    with assert_setup_component(1):
-        await 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')
-    await asyncio.sleep(0.1)
-
-    state = hass.states.get('sensor.connected_clients')
-    assert state.state == '1'
-
-    log.debug('WS %s: %s', id(log), 'Connected')
-    await asyncio.sleep(0.1)
-
-    state = hass.states.get('sensor.connected_clients')
-    assert state.state == '2'
-
-    log.debug('WS %s: %s', id(log), 'Closed connection')
-    await asyncio.sleep(0.1)
-
-    state = hass.states.get('sensor.connected_clients')
-    assert state.state == '1'

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'