Store runtime data inside the config entry in Shelly (#116763)

This commit is contained in:
Michael 2024-05-04 12:41:25 +02:00 committed by GitHub
parent e5543e3b95
commit 8238cd9f22
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 93 additions and 108 deletions

View file

@ -9,7 +9,6 @@ from typing import Any, cast
from aioshelly.block_device import Block
from aioshelly.exceptions import DeviceConnectionError, InvalidAuthError, RpcCallError
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant, State, callback
from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC, DeviceInfo
@ -24,7 +23,7 @@ from homeassistant.helpers.typing import StateType
from homeassistant.helpers.update_coordinator import CoordinatorEntity
from .const import CONF_SLEEP_PERIOD, LOGGER
from .coordinator import ShellyBlockCoordinator, ShellyRpcCoordinator, get_entry_data
from .coordinator import ShellyBlockCoordinator, ShellyConfigEntry, ShellyRpcCoordinator
from .utils import (
async_remove_shelly_entity,
get_block_entity_name,
@ -36,13 +35,13 @@ from .utils import (
@callback
def async_setup_entry_attribute_entities(
hass: HomeAssistant,
config_entry: ConfigEntry,
config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback,
sensors: Mapping[tuple[str, str], BlockEntityDescription],
sensor_class: Callable,
) -> None:
"""Set up entities for attributes."""
coordinator = get_entry_data(hass)[config_entry.entry_id].block
coordinator = config_entry.runtime_data.block
assert coordinator
if coordinator.device.initialized:
async_setup_block_attribute_entities(
@ -104,7 +103,7 @@ def async_setup_block_attribute_entities(
@callback
def async_restore_block_attribute_entities(
hass: HomeAssistant,
config_entry: ConfigEntry,
config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback,
coordinator: ShellyBlockCoordinator,
sensors: Mapping[tuple[str, str], BlockEntityDescription],
@ -139,13 +138,13 @@ def async_restore_block_attribute_entities(
@callback
def async_setup_entry_rpc(
hass: HomeAssistant,
config_entry: ConfigEntry,
config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback,
sensors: Mapping[str, RpcEntityDescription],
sensor_class: Callable,
) -> None:
"""Set up entities for RPC sensors."""
coordinator = get_entry_data(hass)[config_entry.entry_id].rpc
coordinator = config_entry.runtime_data.rpc
assert coordinator
if coordinator.device.initialized:
@ -161,18 +160,18 @@ def async_setup_entry_rpc(
@callback
def async_setup_rpc_attribute_entities(
hass: HomeAssistant,
config_entry: ConfigEntry,
config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback,
sensors: Mapping[str, RpcEntityDescription],
sensor_class: Callable,
) -> None:
"""Set up entities for RPC attributes."""
coordinator = get_entry_data(hass)[config_entry.entry_id].rpc
coordinator = config_entry.runtime_data.rpc
assert coordinator
polling_coordinator = None
if not (sleep_period := config_entry.data[CONF_SLEEP_PERIOD]):
polling_coordinator = get_entry_data(hass)[config_entry.entry_id].rpc_poll
polling_coordinator = config_entry.runtime_data.rpc_poll
assert polling_coordinator
entities = []
@ -213,7 +212,7 @@ def async_setup_rpc_attribute_entities(
@callback
def async_restore_rpc_attribute_entities(
hass: HomeAssistant,
config_entry: ConfigEntry,
config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback,
coordinator: ShellyRpcCoordinator,
sensors: Mapping[str, RpcEntityDescription],
@ -248,13 +247,13 @@ def async_restore_rpc_attribute_entities(
@callback
def async_setup_entry_rest(
hass: HomeAssistant,
config_entry: ConfigEntry,
config_entry: ShellyConfigEntry,
async_add_entities: AddEntitiesCallback,
sensors: Mapping[str, RestEntityDescription],
sensor_class: Callable,
) -> None:
"""Set up entities for REST sensors."""
coordinator = get_entry_data(hass)[config_entry.entry_id].rest
coordinator = config_entry.runtime_data.rest
assert coordinator
async_add_entities(