core/homeassistant/components/websocket_api/sensor.py

"""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()