Refactor Shelly wrapper to coordinator (#79628)

This commit is contained in:
Shay Levy 2022-10-05 15:39:58 +03:00 committed by GitHub
parent 4d3d22320f
commit 22c68b95bf
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
14 changed files with 336 additions and 321 deletions

View file

@ -20,12 +20,12 @@ from homeassistant.components.climate import (
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import ATTR_TEMPERATURE, TEMP_CELSIUS
from homeassistant.core import HomeAssistant, State, callback
from homeassistant.helpers import device_registry, entity_registry, update_coordinator
from homeassistant.helpers import device_registry, entity_registry
from homeassistant.helpers.entity import DeviceInfo
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.restore_state import RestoreEntity
from homeassistant.helpers.update_coordinator import CoordinatorEntity
from . import BlockDeviceWrapper
from .const import (
AIOSHELLY_DEVICE_TIMEOUT_SEC,
BLOCK,
@ -34,6 +34,7 @@ from .const import (
LOGGER,
SHTRV_01_TEMPERATURE_SETTINGS,
)
from .coordinator import ShellyBlockCoordinator
from .utils import get_device_entry_gen
@ -47,37 +48,41 @@ async def async_setup_entry(
if get_device_entry_gen(config_entry) == 2:
return
wrapper: BlockDeviceWrapper = hass.data[DOMAIN][DATA_CONFIG_ENTRY][
coordinator: ShellyBlockCoordinator = hass.data[DOMAIN][DATA_CONFIG_ENTRY][
config_entry.entry_id
][BLOCK]
if wrapper.device.initialized:
async_setup_climate_entities(async_add_entities, wrapper)
if coordinator.device.initialized:
async_setup_climate_entities(async_add_entities, coordinator)
else:
async_restore_climate_entities(hass, config_entry, async_add_entities, wrapper)
async_restore_climate_entities(
hass, config_entry, async_add_entities, coordinator
)
@callback
def async_setup_climate_entities(
async_add_entities: AddEntitiesCallback,
wrapper: BlockDeviceWrapper,
coordinator: ShellyBlockCoordinator,
) -> None:
"""Set up online climate devices."""
device_block: Block | None = None
sensor_block: Block | None = None
assert wrapper.device.blocks
assert coordinator.device.blocks
for block in wrapper.device.blocks:
for block in coordinator.device.blocks:
if block.type == "device":
device_block = block
if hasattr(block, "targetTemp"):
sensor_block = block
if sensor_block and device_block:
LOGGER.debug("Setup online climate device %s", wrapper.name)
async_add_entities([BlockSleepingClimate(wrapper, sensor_block, device_block)])
LOGGER.debug("Setup online climate device %s", coordinator.name)
async_add_entities(
[BlockSleepingClimate(coordinator, sensor_block, device_block)]
)
@callback
@ -85,7 +90,7 @@ def async_restore_climate_entities(
hass: HomeAssistant,
config_entry: ConfigEntry,
async_add_entities: AddEntitiesCallback,
wrapper: BlockDeviceWrapper,
coordinator: ShellyBlockCoordinator,
) -> None:
"""Restore sleeping climate devices."""
@ -99,16 +104,14 @@ def async_restore_climate_entities(
if entry.domain != CLIMATE_DOMAIN:
continue
LOGGER.debug("Setup sleeping climate device %s", wrapper.name)
LOGGER.debug("Setup sleeping climate device %s", coordinator.name)
LOGGER.debug("Found entry %s [%s]", entry.original_name, entry.domain)
async_add_entities([BlockSleepingClimate(wrapper, None, None, entry)])
async_add_entities([BlockSleepingClimate(coordinator, None, None, entry)])
break
class BlockSleepingClimate(
update_coordinator.CoordinatorEntity,
RestoreEntity,
ClimateEntity,
CoordinatorEntity[ShellyBlockCoordinator], RestoreEntity, ClimateEntity
):
"""Representation of a Shelly climate device."""
@ -124,16 +127,14 @@ class BlockSleepingClimate(
def __init__(
self,
wrapper: BlockDeviceWrapper,
coordinator: ShellyBlockCoordinator,
sensor_block: Block | None,
device_block: Block | None,
entry: entity_registry.RegistryEntry | None = None,
) -> None:
"""Initialize climate."""
super().__init__(coordinator)
super().__init__(wrapper)
self.wrapper = wrapper
self.block: Block | None = sensor_block
self.control_result: dict[str, Any] | None = None
self.device_block: Block | None = device_block
@ -142,11 +143,11 @@ class BlockSleepingClimate(
self._preset_modes: list[str] = []
if self.block is not None and self.device_block is not None:
self._unique_id = f"{self.wrapper.mac}-{self.block.description}"
self._unique_id = f"{self.coordinator.mac}-{self.block.description}"
assert self.block.channel
self._preset_modes = [
PRESET_NONE,
*wrapper.device.settings["thermostats"][int(self.block.channel)][
*coordinator.device.settings["thermostats"][int(self.block.channel)][
"schedule_profile_names"
],
]
@ -163,7 +164,7 @@ class BlockSleepingClimate(
@property
def name(self) -> str:
"""Name of entity."""
return self.wrapper.name
return self.coordinator.name
@property
def target_temperature(self) -> float | None:
@ -184,7 +185,7 @@ class BlockSleepingClimate(
"""Device availability."""
if self.device_block is not None:
return not cast(bool, self.device_block.valveError)
return self.wrapper.last_update_success
return self.coordinator.last_update_success
@property
def hvac_mode(self) -> HVACMode:
@ -229,7 +230,9 @@ class BlockSleepingClimate(
def device_info(self) -> DeviceInfo:
"""Device info."""
return {
"connections": {(device_registry.CONNECTION_NETWORK_MAC, self.wrapper.mac)}
"connections": {
(device_registry.CONNECTION_NETWORK_MAC, self.coordinator.mac)
}
}
def _check_is_off(self) -> bool:
@ -244,7 +247,7 @@ class BlockSleepingClimate(
LOGGER.debug("Setting state for entity %s, state: %s", self.name, kwargs)
try:
async with async_timeout.timeout(AIOSHELLY_DEVICE_TIMEOUT_SEC):
return await self.wrapper.device.http_request(
return await self.coordinator.device.http_request(
"get", f"thermostat/{self._channel}", kwargs
)
except (asyncio.TimeoutError, OSError) as err:
@ -254,7 +257,7 @@ class BlockSleepingClimate(
kwargs,
repr(err),
)
self.wrapper.last_update_success = False
self.coordinator.last_update_success = False
return None
async def async_set_temperature(self, **kwargs: Any) -> None:
@ -302,13 +305,13 @@ class BlockSleepingClimate(
@callback
def _handle_coordinator_update(self) -> None:
"""Handle device update."""
if not self.wrapper.device.initialized:
if not self.coordinator.device.initialized:
self.async_write_ha_state()
return
assert self.wrapper.device.blocks
assert self.coordinator.device.blocks
for block in self.wrapper.device.blocks:
for block in self.coordinator.device.blocks:
if block.type == "device":
self.device_block = block
if hasattr(block, "targetTemp"):
@ -322,11 +325,11 @@ class BlockSleepingClimate(
try:
self._preset_modes = [
PRESET_NONE,
*self.wrapper.device.settings["thermostats"][
*self.coordinator.device.settings["thermostats"][
int(self.block.channel)
]["schedule_profile_names"],
]
except AuthRequired:
self.wrapper.entry.async_start_reauth(self.hass)
self.coordinator.entry.async_start_reauth(self.hass)
else:
self.async_write_ha_state()