homeassistant-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."""

    _attr_name = "Connected clients"
    _attr_native_unit_of_measurement = "clients"

    def __init__(self) -> None:
        """Initialize the API count."""
        self._attr_native_value = 0

    async def async_added_to_hass(self) -> None:
        """Handle addition 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
            )
        )

    @callback
    def _update_count(self) -> None:
        self._attr_native_value = self.hass.data.get(DATA_CONNECTIONS, 0)
        self.async_write_ha_state()