Move notion base entity to separate module (#126499)

This commit is contained in:
epenet 2024-09-23 12:45:30 +02:00 committed by GitHub
parent 438cbc99b1
commit 43322bc3d9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 127 additions and 121 deletions

View File

@ -6,18 +6,14 @@ from datetime import timedelta
from typing import Any
from uuid import UUID
from aionotion.bridge.models import Bridge
from aionotion.errors import InvalidCredentialsError, NotionError
from aionotion.listener.models import Listener, ListenerKind
from aionotion.listener.models import ListenerKind
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_PASSWORD, CONF_USERNAME, Platform
from homeassistant.core import HomeAssistant, callback
from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
from homeassistant.helpers import device_registry as dr, entity_registry as er
from homeassistant.helpers.device_registry import DeviceInfo
from homeassistant.helpers.entity import EntityDescription
from homeassistant.helpers.update_coordinator import CoordinatorEntity
from homeassistant.helpers import entity_registry as er
from .const import (
CONF_REFRESH_TOKEN,
@ -168,102 +164,3 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
hass.data[DOMAIN].pop(entry.entry_id)
return unload_ok
class NotionEntity(CoordinatorEntity[NotionDataUpdateCoordinator]):
"""Define a base Notion entity."""
_attr_has_entity_name = True
def __init__(
self,
coordinator: NotionDataUpdateCoordinator,
listener_id: str,
sensor_id: str,
bridge_id: int,
description: EntityDescription,
) -> None:
"""Initialize the entity."""
super().__init__(coordinator)
sensor = self.coordinator.data.sensors[sensor_id]
self._attr_device_info = DeviceInfo(
identifiers={(DOMAIN, sensor.hardware_id)},
manufacturer="Silicon Labs",
model=str(sensor.hardware_revision),
name=str(sensor.name).capitalize(),
sw_version=sensor.firmware_version,
)
if bridge := self._async_get_bridge(bridge_id):
self._attr_device_info["via_device"] = (DOMAIN, bridge.hardware_id)
self._attr_extra_state_attributes = {}
self._attr_unique_id = listener_id
self._bridge_id = bridge_id
self._listener_id = listener_id
self._sensor_id = sensor_id
self.entity_description = description
@property
def available(self) -> bool:
"""Return True if entity is available."""
return (
self.coordinator.last_update_success
and self._listener_id in self.coordinator.data.listeners
)
@property
def listener(self) -> Listener:
"""Return the listener related to this entity."""
return self.coordinator.data.listeners[self._listener_id]
@callback
def _async_get_bridge(self, bridge_id: int) -> Bridge | None:
"""Get a bridge by ID (if it exists)."""
if (bridge := self.coordinator.data.bridges.get(bridge_id)) is None:
LOGGER.debug("Entity references a non-existent bridge ID: %s", bridge_id)
return None
return bridge
@callback
def _async_update_bridge_id(self) -> None:
"""Update the entity's bridge ID if it has changed.
Sensors can move to other bridges based on signal strength, etc.
"""
sensor = self.coordinator.data.sensors[self._sensor_id]
# If the bridge ID hasn't changed, return:
if self._bridge_id == sensor.bridge.id:
return
# If the bridge doesn't exist, return:
if (bridge := self._async_get_bridge(sensor.bridge.id)) is None:
return
self._bridge_id = sensor.bridge.id
device_registry = dr.async_get(self.hass)
this_device = device_registry.async_get_device(
identifiers={(DOMAIN, sensor.hardware_id)}
)
bridge = self.coordinator.data.bridges[self._bridge_id]
bridge_device = device_registry.async_get_device(
identifiers={(DOMAIN, bridge.hardware_id)}
)
if not bridge_device or not this_device:
return
device_registry.async_update_device(
this_device.id, via_device_id=bridge_device.id
)
@callback
def _handle_coordinator_update(self) -> None:
"""Respond to a DataUpdateCoordinator update."""
if self._listener_id in self.coordinator.data.listeners:
self._async_update_bridge_id()
super()._handle_coordinator_update()

View File

@ -17,7 +17,6 @@ from homeassistant.const import EntityCategory
from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from . import NotionEntity
from .const import (
DOMAIN,
LOGGER,
@ -32,7 +31,7 @@ from .const import (
SENSOR_WINDOW_HINGED,
)
from .coordinator import NotionDataUpdateCoordinator
from .model import NotionEntityDescription
from .entity import NotionEntity, NotionEntityDescription
@dataclass(frozen=True, kw_only=True)

View File

@ -0,0 +1,123 @@
"""Support for Notion."""
from __future__ import annotations
from dataclasses import dataclass
from aionotion.bridge.models import Bridge
from aionotion.listener.models import Listener, ListenerKind
from homeassistant.core import callback
from homeassistant.helpers import device_registry as dr
from homeassistant.helpers.device_registry import DeviceInfo
from homeassistant.helpers.entity import EntityDescription
from homeassistant.helpers.update_coordinator import CoordinatorEntity
from .const import DOMAIN, LOGGER
from .coordinator import NotionDataUpdateCoordinator
@dataclass(frozen=True, kw_only=True)
class NotionEntityDescription:
"""Define an description for Notion entities."""
listener_kind: ListenerKind
class NotionEntity(CoordinatorEntity[NotionDataUpdateCoordinator]):
"""Define a base Notion entity."""
_attr_has_entity_name = True
def __init__(
self,
coordinator: NotionDataUpdateCoordinator,
listener_id: str,
sensor_id: str,
bridge_id: int,
description: EntityDescription,
) -> None:
"""Initialize the entity."""
super().__init__(coordinator)
sensor = self.coordinator.data.sensors[sensor_id]
self._attr_device_info = DeviceInfo(
identifiers={(DOMAIN, sensor.hardware_id)},
manufacturer="Silicon Labs",
model=str(sensor.hardware_revision),
name=str(sensor.name).capitalize(),
sw_version=sensor.firmware_version,
)
if bridge := self._async_get_bridge(bridge_id):
self._attr_device_info["via_device"] = (DOMAIN, bridge.hardware_id)
self._attr_extra_state_attributes = {}
self._attr_unique_id = listener_id
self._bridge_id = bridge_id
self._listener_id = listener_id
self._sensor_id = sensor_id
self.entity_description = description
@property
def available(self) -> bool:
"""Return True if entity is available."""
return (
self.coordinator.last_update_success
and self._listener_id in self.coordinator.data.listeners
)
@property
def listener(self) -> Listener:
"""Return the listener related to this entity."""
return self.coordinator.data.listeners[self._listener_id]
@callback
def _async_get_bridge(self, bridge_id: int) -> Bridge | None:
"""Get a bridge by ID (if it exists)."""
if (bridge := self.coordinator.data.bridges.get(bridge_id)) is None:
LOGGER.debug("Entity references a non-existent bridge ID: %s", bridge_id)
return None
return bridge
@callback
def _async_update_bridge_id(self) -> None:
"""Update the entity's bridge ID if it has changed.
Sensors can move to other bridges based on signal strength, etc.
"""
sensor = self.coordinator.data.sensors[self._sensor_id]
# If the bridge ID hasn't changed, return:
if self._bridge_id == sensor.bridge.id:
return
# If the bridge doesn't exist, return:
if (bridge := self._async_get_bridge(sensor.bridge.id)) is None:
return
self._bridge_id = sensor.bridge.id
device_registry = dr.async_get(self.hass)
this_device = device_registry.async_get_device(
identifiers={(DOMAIN, sensor.hardware_id)}
)
bridge = self.coordinator.data.bridges[self._bridge_id]
bridge_device = device_registry.async_get_device(
identifiers={(DOMAIN, bridge.hardware_id)}
)
if not bridge_device or not this_device:
return
device_registry.async_update_device(
this_device.id, via_device_id=bridge_device.id
)
@callback
def _handle_coordinator_update(self) -> None:
"""Respond to a DataUpdateCoordinator update."""
if self._listener_id in self.coordinator.data.listeners:
self._async_update_bridge_id()
super()._handle_coordinator_update()

View File

@ -1,12 +0,0 @@
"""Define Notion model mixins."""
from dataclasses import dataclass
from aionotion.listener.models import ListenerKind
@dataclass(frozen=True, kw_only=True)
class NotionEntityDescription:
"""Define an description for Notion entities."""
listener_kind: ListenerKind

View File

@ -15,10 +15,9 @@ from homeassistant.const import UnitOfTemperature
from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from . import NotionEntity
from .const import DOMAIN, SENSOR_MOLD, SENSOR_TEMPERATURE
from .coordinator import NotionDataUpdateCoordinator
from .model import NotionEntityDescription
from .entity import NotionEntity, NotionEntityDescription
@dataclass(frozen=True, kw_only=True)