core/homeassistant/components/overkiz/water_heater.py

"""Support for Overkiz water heater devices."""

from __future__ import annotations

from homeassistant.config_entries import ConfigEntry
from homeassistant.const import Platform
from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity_platform import AddEntitiesCallback

from . import HomeAssistantOverkizData
from .const import DOMAIN
from .entity import OverkizEntity
from .water_heater_entities import (
    CONTROLLABLE_NAME_TO_WATER_HEATER_ENTITY,
    WIDGET_TO_WATER_HEATER_ENTITY,
)


async def async_setup_entry(
    hass: HomeAssistant,
    entry: ConfigEntry,
    async_add_entities: AddEntitiesCallback,
) -> None:
    """Set up the Overkiz DHW from a config entry."""
    data: HomeAssistantOverkizData = hass.data[DOMAIN][entry.entry_id]
    entities: list[OverkizEntity] = []

    for device in data.platforms[Platform.WATER_HEATER]:
        if device.controllable_name in CONTROLLABLE_NAME_TO_WATER_HEATER_ENTITY:
            entities.append(
                CONTROLLABLE_NAME_TO_WATER_HEATER_ENTITY[device.controllable_name](
                    device.device_url, data.coordinator
                )
            )
        elif device.widget in WIDGET_TO_WATER_HEATER_ENTITY:
            entities.append(
                WIDGET_TO_WATER_HEATER_ENTITY[device.widget](
                    device.device_url, data.coordinator
                )
            )

    async_add_entities(entities)