68 lines
1.9 KiB
Python
68 lines
1.9 KiB
Python
"""Entity to track connections to websocket API."""
|
|
from __future__ import annotations
|
|
|
|
from homeassistant.components.sensor import SensorEntity
|
|
from homeassistant.core import HomeAssistant, callback
|
|
from homeassistant.helpers.dispatcher import async_dispatcher_connect
|
|
from homeassistant.helpers.entity_platform import AddEntitiesCallback
|
|
from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
|
|
|
|
from .const import (
|
|
DATA_CONNECTIONS,
|
|
SIGNAL_WEBSOCKET_CONNECTED,
|
|
SIGNAL_WEBSOCKET_DISCONNECTED,
|
|
)
|
|
|
|
|
|
async def async_setup_platform(
|
|
hass: HomeAssistant,
|
|
config: ConfigType,
|
|
async_add_entities: AddEntitiesCallback,
|
|
discovery_info: DiscoveryInfoType | None = None,
|
|
) -> None:
|
|
"""Set up the API streams platform."""
|
|
entity = APICount()
|
|
|
|
async_add_entities([entity])
|
|
|
|
|
|
class APICount(SensorEntity):
|
|
"""Entity to represent how many people are connected to the stream API."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Initialize the API count."""
|
|
self.count = 0
|
|
|
|
async def async_added_to_hass(self) -> None:
|
|
"""Added to hass."""
|
|
self.async_on_remove(
|
|
async_dispatcher_connect(
|
|
self.hass, SIGNAL_WEBSOCKET_CONNECTED, self._update_count
|
|
)
|
|
)
|
|
self.async_on_remove(
|
|
async_dispatcher_connect(
|
|
self.hass, SIGNAL_WEBSOCKET_DISCONNECTED, self._update_count
|
|
)
|
|
)
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
"""Return name of entity."""
|
|
return "Connected clients"
|
|
|
|
@property
|
|
def native_value(self) -> int:
|
|
"""Return current API count."""
|
|
return self.count
|
|
|
|
@property
|
|
def native_unit_of_measurement(self) -> str:
|
|
"""Return the unit of measurement."""
|
|
return "clients"
|
|
|
|
@callback
|
|
def _update_count(self) -> None:
|
|
self.count = self.hass.data.get(DATA_CONNECTIONS, 0)
|
|
self.async_write_ha_state()
|