Add Shelly support for sleeping Gen2 devices (#79889)

This commit is contained in:
Shay Levy 2022-10-18 22:42:22 +03:00 committed by GitHub
parent 98ca28ab1c
commit 8e9457d808
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
11 changed files with 337 additions and 54 deletions

View file

@ -39,7 +39,13 @@ from .coordinator import (
ShellyRpcPollingCoordinator,
get_entry_data,
)
from .utils import get_block_device_sleep_period, get_coap_context, get_device_entry_gen
from .utils import (
get_block_device_sleep_period,
get_coap_context,
get_device_entry_gen,
get_rpc_device_sleep_period,
get_ws_context,
)
BLOCK_PLATFORMS: Final = [
Platform.BINARY_SENSOR,
@ -65,7 +71,10 @@ RPC_PLATFORMS: Final = [
Platform.SWITCH,
Platform.UPDATE,
]
RPC_SLEEPING_PLATFORMS: Final = [
Platform.BINARY_SENSOR,
Platform.SENSOR,
]
COAP_SCHEMA: Final = vol.Schema(
{
@ -215,26 +224,87 @@ async def _async_setup_rpc_entry(hass: HomeAssistant, entry: ConfigEntry) -> boo
entry.data.get(CONF_PASSWORD),
)
LOGGER.debug("Setting up online RPC device %s", entry.title)
try:
async with async_timeout.timeout(AIOSHELLY_DEVICE_TIMEOUT_SEC):
device = await RpcDevice.create(
aiohttp_client.async_get_clientsession(hass), options
)
except asyncio.TimeoutError as err:
raise ConfigEntryNotReady(str(err) or "Timeout during device setup") from err
except OSError as err:
raise ConfigEntryNotReady(str(err) or "Error during device setup") from err
except (AuthRequired, InvalidAuthError) as err:
raise ConfigEntryAuthFailed from err
ws_context = await get_ws_context(hass)
device = await RpcDevice.create(
aiohttp_client.async_get_clientsession(hass),
ws_context,
options,
False,
)
dev_reg = device_registry.async_get(hass)
device_entry = None
if entry.unique_id is not None:
device_entry = dev_reg.async_get_device(
identifiers=set(),
connections={
(
device_registry.CONNECTION_NETWORK_MAC,
device_registry.format_mac(entry.unique_id),
)
},
)
if device_entry and entry.entry_id not in device_entry.config_entries:
device_entry = None
sleep_period = entry.data.get(CONF_SLEEP_PERIOD)
shelly_entry_data = get_entry_data(hass)[entry.entry_id]
shelly_entry_data.rpc = ShellyRpcCoordinator(hass, entry, device)
shelly_entry_data.rpc.async_setup()
shelly_entry_data.rpc_poll = ShellyRpcPollingCoordinator(hass, entry, device)
@callback
def _async_rpc_device_setup() -> None:
"""Set up a RPC based device that is online."""
shelly_entry_data.rpc = ShellyRpcCoordinator(hass, entry, device)
shelly_entry_data.rpc.async_setup()
hass.config_entries.async_setup_platforms(entry, RPC_PLATFORMS)
platforms = RPC_SLEEPING_PLATFORMS
if not entry.data.get(CONF_SLEEP_PERIOD):
shelly_entry_data.rpc_poll = ShellyRpcPollingCoordinator(
hass, entry, device
)
platforms = RPC_PLATFORMS
hass.config_entries.async_setup_platforms(entry, platforms)
@callback
def _async_device_online(_: Any) -> None:
LOGGER.debug("Device %s is online, resuming setup", entry.title)
shelly_entry_data.device = None
if sleep_period is None:
data = {**entry.data}
data[CONF_SLEEP_PERIOD] = get_rpc_device_sleep_period(device.config)
hass.config_entries.async_update_entry(entry, data=data)
_async_rpc_device_setup()
if sleep_period == 0:
# Not a sleeping device, finish setup
LOGGER.debug("Setting up online RPC device %s", entry.title)
try:
async with async_timeout.timeout(AIOSHELLY_DEVICE_TIMEOUT_SEC):
await device.initialize()
except asyncio.TimeoutError as err:
raise ConfigEntryNotReady(
str(err) or "Timeout during device setup"
) from err
except OSError as err:
raise ConfigEntryNotReady(str(err) or "Error during device setup") from err
except (AuthRequired, InvalidAuthError) as err:
raise ConfigEntryAuthFailed from err
_async_rpc_device_setup()
elif sleep_period is None or device_entry is None:
# Need to get sleep info or first time sleeping device setup, wait for device
shelly_entry_data.device = device
LOGGER.debug(
"Setup for device %s will resume when device is online", entry.title
)
device.subscribe_updates(_async_device_online)
else:
# Restore sensors for sleeping device
LOGGER.debug("Setting up offline block device %s", entry.title)
_async_rpc_device_setup()
return True
@ -243,9 +313,18 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Unload a config entry."""
shelly_entry_data = get_entry_data(hass)[entry.entry_id]
if shelly_entry_data.device is not None:
# If device is present, block/rpc coordinator is not setup yet
shelly_entry_data.device.shutdown()
return True
platforms = RPC_SLEEPING_PLATFORMS
if not entry.data.get(CONF_SLEEP_PERIOD):
platforms = RPC_PLATFORMS
if get_device_entry_gen(entry) == 2:
if unload_ok := await hass.config_entries.async_unload_platforms(
entry, RPC_PLATFORMS
entry, platforms
):
if shelly_entry_data.rpc:
await shelly_entry_data.rpc.shutdown()
@ -253,11 +332,6 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
return unload_ok
if shelly_entry_data.device is not None:
# If device is present, block coordinator is not setup yet
shelly_entry_data.device.shutdown()
return True
platforms = BLOCK_SLEEPING_PLATFORMS
if not entry.data.get(CONF_SLEEP_PERIOD):

View file

@ -25,6 +25,7 @@ from .entity import (
ShellyRestAttributeEntity,
ShellyRpcAttributeEntity,
ShellySleepingBlockAttributeEntity,
ShellySleepingRpcAttributeEntity,
async_setup_entry_attribute_entities,
async_setup_entry_rest,
async_setup_entry_rpc,
@ -209,9 +210,19 @@ async def async_setup_entry(
) -> None:
"""Set up sensors for device."""
if get_device_entry_gen(config_entry) == 2:
return async_setup_entry_rpc(
hass, config_entry, async_add_entities, RPC_SENSORS, RpcBinarySensor
)
if config_entry.data[CONF_SLEEP_PERIOD]:
async_setup_entry_rpc(
hass,
config_entry,
async_add_entities,
RPC_SENSORS,
RpcSleepingBinarySensor,
)
else:
async_setup_entry_rpc(
hass, config_entry, async_add_entities, RPC_SENSORS, RpcBinarySensor
)
return
if config_entry.data[CONF_SLEEP_PERIOD]:
async_setup_entry_attribute_entities(
@ -289,3 +300,17 @@ class BlockSleepingBinarySensor(ShellySleepingBlockAttributeEntity, BinarySensor
return bool(self.attribute_value)
return self.last_state == STATE_ON
class RpcSleepingBinarySensor(ShellySleepingRpcAttributeEntity, BinarySensorEntity):
"""Represent a RPC sleeping binary sensor entity."""
entity_description: RpcBinarySensorDescription
@property
def is_on(self) -> bool | None:
"""Return true if RPC sensor state is on."""
if self.coordinator.device.initialized:
return bool(self.attribute_value)
return self.last_state == STATE_ON

View file

@ -29,6 +29,8 @@ from .utils import (
get_info_gen,
get_model_name,
get_rpc_device_name,
get_rpc_device_sleep_period,
get_ws_context,
)
HOST_SCHEMA: Final = vol.Schema({vol.Required(CONF_HOST): str})
@ -54,8 +56,10 @@ async def validate_input(
async with async_timeout.timeout(AIOSHELLY_DEVICE_TIMEOUT_SEC):
if get_info_gen(info) == 2:
ws_context = await get_ws_context(hass)
rpc_device = await RpcDevice.create(
aiohttp_client.async_get_clientsession(hass),
ws_context,
options,
)
await rpc_device.shutdown()
@ -63,7 +67,7 @@ async def validate_input(
return {
"title": get_rpc_device_name(rpc_device),
CONF_SLEEP_PERIOD: 0,
CONF_SLEEP_PERIOD: get_rpc_device_sleep_period(rpc_device.config),
"model": rpc_device.shelly.get("model"),
"gen": 2,
}

View file

@ -53,7 +53,7 @@ class ShellyEntryData:
"""Class for sharing data within a given config entry."""
block: ShellyBlockCoordinator | None = None
device: BlockDevice | None = None
device: BlockDevice | RpcDevice | None = None
rest: ShellyRestCoordinator | None = None
rpc: ShellyRpcCoordinator | None = None
rpc_poll: ShellyRpcPollingCoordinator | None = None
@ -353,12 +353,16 @@ class ShellyRpcCoordinator(DataUpdateCoordinator):
"""Initialize the Shelly RPC device coordinator."""
self.device_id: str | None = None
if sleep_period := entry.data[CONF_SLEEP_PERIOD]:
update_interval = SLEEP_PERIOD_MULTIPLIER * sleep_period
else:
update_interval = RPC_RECONNECT_INTERVAL
device_name = get_rpc_device_name(device) if device.initialized else entry.title
super().__init__(
hass,
LOGGER,
name=device_name,
update_interval=timedelta(seconds=RPC_RECONNECT_INTERVAL),
update_interval=timedelta(seconds=update_interval),
)
self.entry = entry
self.device = device
@ -424,6 +428,11 @@ class ShellyRpcCoordinator(DataUpdateCoordinator):
async def _async_update_data(self) -> None:
"""Fetch data."""
if sleep_period := self.entry.data.get(CONF_SLEEP_PERIOD):
# Sleeping device, no point polling it, just mark it unavailable
raise UpdateFailed(
f"Sleeping device did not update within {sleep_period} seconds interval"
)
if self.device.connected:
return

View file

@ -14,11 +14,12 @@ from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers import device_registry, entity, entity_registry
from homeassistant.helpers.entity import DeviceInfo, EntityDescription
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.entity_registry import RegistryEntry
from homeassistant.helpers.restore_state import RestoreEntity
from homeassistant.helpers.typing import StateType
from homeassistant.helpers.update_coordinator import CoordinatorEntity
from .const import AIOSHELLY_DEVICE_TIMEOUT_SEC, LOGGER
from .const import AIOSHELLY_DEVICE_TIMEOUT_SEC, CONF_SLEEP_PERIOD, LOGGER
from .coordinator import (
ShellyBlockCoordinator,
ShellyRpcCoordinator,
@ -40,9 +41,7 @@ def async_setup_entry_attribute_entities(
async_add_entities: AddEntitiesCallback,
sensors: Mapping[tuple[str, str], BlockEntityDescription],
sensor_class: Callable,
description_class: Callable[
[entity_registry.RegistryEntry], BlockEntityDescription
],
description_class: Callable[[RegistryEntry], BlockEntityDescription],
) -> None:
"""Set up entities for attributes."""
coordinator = get_entry_data(hass)[config_entry.entry_id].block
@ -115,9 +114,7 @@ def async_restore_block_attribute_entities(
coordinator: ShellyBlockCoordinator,
sensors: Mapping[tuple[str, str], BlockEntityDescription],
sensor_class: Callable,
description_class: Callable[
[entity_registry.RegistryEntry], BlockEntityDescription
],
description_class: Callable[[RegistryEntry], BlockEntityDescription],
) -> None:
"""Restore block attributes entities."""
entities = []
@ -154,11 +151,35 @@ def async_setup_entry_rpc(
sensors: Mapping[str, RpcEntityDescription],
sensor_class: Callable,
) -> None:
"""Set up entities for REST sensors."""
"""Set up entities for RPC sensors."""
coordinator = get_entry_data(hass)[config_entry.entry_id].rpc
assert coordinator
polling_coordinator = get_entry_data(hass)[config_entry.entry_id].rpc_poll
assert polling_coordinator
if coordinator.device.initialized:
async_setup_rpc_attribute_entities(
hass, config_entry, async_add_entities, sensors, sensor_class
)
else:
async_restore_rpc_attribute_entities(
hass, config_entry, async_add_entities, coordinator, sensors, sensor_class
)
@callback
def async_setup_rpc_attribute_entities(
hass: HomeAssistant,
config_entry: ConfigEntry,
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
assert coordinator
if not (sleep_period := config_entry.data[CONF_SLEEP_PERIOD]):
polling_coordinator = get_entry_data(hass)[config_entry.entry_id].rpc_poll
assert polling_coordinator
entities = []
for sensor_id in sensors:
@ -183,13 +204,52 @@ def async_setup_entry_rpc(
async_remove_shelly_entity(hass, domain, unique_id)
else:
if description.use_polling_coordinator:
entities.append(
sensor_class(polling_coordinator, key, sensor_id, description)
)
if not sleep_period:
entities.append(
sensor_class(
polling_coordinator, key, sensor_id, description
)
)
else:
entities.append(
sensor_class(coordinator, key, sensor_id, description)
)
if not entities:
return
async_add_entities(entities)
@callback
def async_restore_rpc_attribute_entities(
hass: HomeAssistant,
config_entry: ConfigEntry,
async_add_entities: AddEntitiesCallback,
coordinator: ShellyRpcCoordinator,
sensors: Mapping[str, RpcEntityDescription],
sensor_class: Callable,
) -> None:
"""Restore block attributes entities."""
entities = []
ent_reg = entity_registry.async_get(hass)
entries = entity_registry.async_entries_for_config_entry(
ent_reg, config_entry.entry_id
)
domain = sensor_class.__module__.split(".")[-1]
for entry in entries:
if entry.domain != domain:
continue
key = entry.unique_id.split("-")[-2]
attribute = entry.unique_id.split("-")[-1]
if description := sensors.get(attribute):
entities.append(
sensor_class(coordinator, key, attribute, description, entry)
)
if not entities:
return
@ -336,7 +396,7 @@ class ShellyRpcEntity(entity.Entity):
@property
def available(self) -> bool:
"""Available."""
return self.coordinator.device.connected
return self.coordinator.last_update_success
@property
def status(self) -> dict:
@ -552,7 +612,7 @@ class ShellySleepingBlockAttributeEntity(ShellyBlockAttributeEntity, RestoreEnti
block: Block | None,
attribute: str,
description: BlockEntityDescription,
entry: entity_registry.RegistryEntry | None = None,
entry: RegistryEntry | None = None,
sensors: Mapping[tuple[str, str], BlockEntityDescription] | None = None,
) -> None:
"""Initialize the sleeping sensor."""
@ -621,3 +681,50 @@ class ShellySleepingBlockAttributeEntity(ShellyBlockAttributeEntity, RestoreEnti
LOGGER.debug("Entity %s attached to block", self.name)
super()._update_callback()
return
class ShellySleepingRpcAttributeEntity(ShellyRpcAttributeEntity, RestoreEntity):
"""Helper class to represent a sleeping rpc attribute."""
entity_description: RpcEntityDescription
# pylint: disable=super-init-not-called
def __init__(
self,
coordinator: ShellyRpcCoordinator,
key: str,
attribute: str,
description: RpcEntityDescription,
entry: RegistryEntry | None = None,
) -> None:
"""Initialize the sleeping sensor."""
self.last_state: StateType = None
self.coordinator = coordinator
self.key = key
self.attribute = attribute
self.entity_description = description
self._attr_should_poll = False
self._attr_device_info = DeviceInfo(
connections={(device_registry.CONNECTION_NETWORK_MAC, coordinator.mac)}
)
self._attr_unique_id = (
self._attr_unique_id
) = f"{coordinator.mac}-{key}-{attribute}"
self._last_value = None
if coordinator.device.initialized:
self._attr_name = get_rpc_entity_name(
coordinator.device, key, description.name
)
elif entry is not None:
self._attr_name = cast(str, entry.original_name)
async def async_added_to_hass(self) -> None:
"""Handle entity which will be added."""
await super().async_added_to_hass()
last_state = await self.async_get_last_state()
if last_state is not None:
self.last_state = last_state.state

View file

@ -3,7 +3,8 @@
"name": "Shelly",
"config_flow": true,
"documentation": "https://www.home-assistant.io/integrations/shelly",
"requirements": ["aioshelly==2.0.2"],
"requirements": ["aioshelly==3.0.0"],
"dependencies": ["http"],
"zeroconf": [
{
"type": "_http._tcp.local.",

View file

@ -42,6 +42,7 @@ from .entity import (
ShellyRestAttributeEntity,
ShellyRpcAttributeEntity,
ShellySleepingBlockAttributeEntity,
ShellySleepingRpcAttributeEntity,
async_setup_entry_attribute_entities,
async_setup_entry_rest,
async_setup_entry_rpc,
@ -451,9 +452,19 @@ async def async_setup_entry(
) -> None:
"""Set up sensors for device."""
if get_device_entry_gen(config_entry) == 2:
return async_setup_entry_rpc(
hass, config_entry, async_add_entities, RPC_SENSORS, RpcSensor
)
if config_entry.data[CONF_SLEEP_PERIOD]:
async_setup_entry_rpc(
hass,
config_entry,
async_add_entities,
RPC_SENSORS,
RpcSleepingSensor,
)
else:
async_setup_entry_rpc(
hass, config_entry, async_add_entities, RPC_SENSORS, RpcSensor
)
return
if config_entry.data[CONF_SLEEP_PERIOD]:
async_setup_entry_attribute_entities(
@ -553,3 +564,17 @@ class BlockSleepingSensor(ShellySleepingBlockAttributeEntity, SensorEntity):
return self.attribute_value
return self.last_state
class RpcSleepingSensor(ShellySleepingRpcAttributeEntity, SensorEntity):
"""Represent a RPC sleeping sensor."""
entity_description: RpcSensorDescription
@property
def native_value(self) -> StateType:
"""Return value of sensor."""
if self.coordinator.device.initialized:
return self.attribute_value
return self.last_state

View file

@ -4,10 +4,12 @@ from __future__ import annotations
from datetime import datetime, timedelta
from typing import Any, cast
from aiohttp.web import Request, WebSocketResponse
from aioshelly.block_device import BLOCK_VALUE_UNIT, COAP, Block, BlockDevice
from aioshelly.const import MODEL_NAMES
from aioshelly.rpc_device import RpcDevice
from aioshelly.rpc_device import RpcDevice, WsServer
from homeassistant.components.http import HomeAssistantView
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import EVENT_HOMEASSISTANT_STOP, TEMP_CELSIUS, TEMP_FAHRENHEIT
from homeassistant.core import HomeAssistant, callback
@ -212,7 +214,7 @@ def get_shbtn_input_triggers() -> list[tuple[str, str]]:
@singleton.singleton("shelly_coap")
async def get_coap_context(hass: HomeAssistant) -> COAP:
"""Get CoAP context to be used in all Shelly devices."""
"""Get CoAP context to be used in all Shelly Gen1 devices."""
context = COAP()
if DOMAIN in hass.data:
port = hass.data[DOMAIN].get(CONF_COAP_PORT, DEFAULT_COAP_PORT)
@ -226,10 +228,33 @@ async def get_coap_context(hass: HomeAssistant) -> COAP:
context.close()
hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, shutdown_listener)
return context
class ShellyReceiver(HomeAssistantView):
"""Handle pushes from Shelly Gen2 devices."""
requires_auth = False
url = "/api/shelly/ws"
name = "api:shelly:ws"
def __init__(self, ws_server: WsServer) -> None:
"""Initialize the Shelly receiver view."""
self._ws_server = ws_server
async def get(self, request: Request) -> WebSocketResponse:
"""Start a get request."""
return await self._ws_server.websocket_handler(request)
@singleton.singleton("shelly_ws_server")
async def get_ws_context(hass: HomeAssistant) -> WsServer:
"""Get websocket server context to be used in all Shelly Gen2 devices."""
ws_server = WsServer()
hass.http.register_view(ShellyReceiver(ws_server))
return ws_server
def get_block_device_sleep_period(settings: dict[str, Any]) -> int:
"""Return the device sleep period in seconds or 0 for non sleeping devices."""
sleep_period = 0
@ -242,6 +267,11 @@ def get_block_device_sleep_period(settings: dict[str, Any]) -> int:
return sleep_period * 60 # minutes to seconds
def get_rpc_device_sleep_period(config: dict[str, Any]) -> int:
"""Return the device sleep period in seconds or 0 for non sleeping devices."""
return cast(int, config["sys"].get("sleep", {}).get("wakeup_period", 0))
def get_info_auth(info: dict[str, Any]) -> bool:
"""Return true if device has authorization enabled."""
return cast(bool, info.get("auth") or info.get("auth_en"))

View file

@ -255,7 +255,7 @@ aiosenseme==0.6.1
aiosenz==1.0.0
# homeassistant.components.shelly
aioshelly==2.0.2
aioshelly==3.0.0
# homeassistant.components.skybell
aioskybell==22.7.0

View file

@ -230,7 +230,7 @@ aiosenseme==0.6.1
aiosenz==1.0.0
# homeassistant.components.shelly
aioshelly==2.0.2
aioshelly==3.0.0
# homeassistant.components.skybell
aioskybell==22.7.0

View file

@ -146,6 +146,13 @@ def mock_coap():
yield
@pytest.fixture(autouse=True)
def mock_ws_server():
"""Mock out ws_server."""
with patch("homeassistant.components.shelly.utils.get_ws_context"):
yield
@pytest.fixture
def device_reg(hass):
"""Return an empty, loaded, registry."""
@ -211,6 +218,7 @@ async def mock_rpc_device():
update=AsyncMock(),
trigger_ota_update=AsyncMock(),
trigger_reboot=AsyncMock(),
initialize=AsyncMock(),
initialized=True,
shutdown=AsyncMock(),
)