diff --git a/homeassistant/components/sensor/api_streams.py b/homeassistant/components/sensor/api_streams.py index 232ff74e862..15cfc200c4d 100644 --- a/homeassistant/components/sensor/api_streams.py +++ b/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]) diff --git a/tests/components/sensor/test_api_streams.py b/tests/components/sensor/test_api_streams.py index 2245a650a1a..2154cc3bd49 100644 --- a/tests/components/sensor/test_api_streams.py +++ b/tests/components/sensor/test_api_streams.py @@ -37,3 +37,37 @@ def test_api_streams(hass): state = hass.states.get('sensor.connected_clients') assert state.state == '1' + + +@asyncio.coroutine +def test_websocket_api(hass): + """Test API streams.""" + log = logging.getLogger('homeassistant.components.websocket_api') + + with assert_setup_component(1): + yield from 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') + yield from hass.async_block_till_done() + + state = hass.states.get('sensor.connected_clients') + assert state.state == '1' + + log.debug('WS %s: %s', id(log), 'Connected') + yield from hass.async_block_till_done() + + state = hass.states.get('sensor.connected_clients') + assert state.state == '2' + + log.debug('WS %s: %s', id(log), 'Closed connection') + yield from hass.async_block_till_done() + + state = hass.states.get('sensor.connected_clients') + assert state.state == '1'