diff --git a/custom_components/sagemcom_fast/__init__.py b/custom_components/sagemcom_fast/__init__.py index 76253ea..81b2c19 100644 --- a/custom_components/sagemcom_fast/__init__.py +++ b/custom_components/sagemcom_fast/__init__.py @@ -1,7 +1,7 @@ """The Sagemcom F@st integration.""" from __future__ import annotations -import asyncio +from dataclasses import dataclass from datetime import timedelta from aiohttp.client_exceptions import ClientError @@ -26,6 +26,7 @@ MaximumSessionCountException, UnauthorizedException, ) +from sagemcom_api.models import DeviceInfo as GatewayDeviceInfo from .const import ( CONF_ENCRYPTION_METHOD, @@ -39,8 +40,16 @@ SERVICE_REBOOT = "reboot" +@dataclass +class HomeAssistantSagemcomFastData: + """SagemcomFast data stored in the Home Assistant data object.""" + + coordinator: SagemcomDataUpdateCoordinator + gateway: GatewayDeviceInfo + + async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry): - """Set up Sagemcom from a config entry.""" + """Set up Sagemcom F@st from a config entry.""" host = entry.data[CONF_HOST] username = entry.data[CONF_USERNAME] @@ -94,10 +103,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry): await coordinator.async_refresh() - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = { - "coordinator": coordinator, - "update_listener": entry.add_update_listener(update_listener), - } + hass.data.setdefault(DOMAIN, {})[entry.entry_id] = HomeAssistantSagemcomFastData( + coordinator=coordinator, gateway=gateway + ) # Create gateway device in Home Assistant device_registry = hass.helpers.device_registry.async_get(hass) @@ -125,22 +133,15 @@ async def async_command_reboot(call): hass, DOMAIN, SERVICE_REBOOT, async_command_reboot ) + entry.async_on_unload(entry.add_update_listener(update_listener)) + return True -async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry): +async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Unload a config entry.""" - unload_ok = all( - await asyncio.gather( - *[ - hass.config_entries.async_forward_entry_unload(entry, component) - for component in PLATFORMS - ] - ) - ) - if unload_ok: - hass.data[DOMAIN][entry.entry_id]["update_listener"]() + if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS): hass.data[DOMAIN].pop(entry.entry_id) return unload_ok diff --git a/custom_components/sagemcom_fast/device_tracker.py b/custom_components/sagemcom_fast/device_tracker.py index de151e7..b514d14 100644 --- a/custom_components/sagemcom_fast/device_tracker.py +++ b/custom_components/sagemcom_fast/device_tracker.py @@ -9,6 +9,7 @@ from homeassistant.helpers.update_coordinator import CoordinatorEntity from sagemcom_api.models import Device +from . import HomeAssistantSagemcomFastData from .const import DOMAIN from .coordinator import SagemcomDataUpdateCoordinator @@ -16,11 +17,11 @@ async def async_setup_entry(hass, config_entry, async_add_entities): """Set up from config entry.""" - coordinator = hass.data[DOMAIN][config_entry.entry_id]["coordinator"] + data: HomeAssistantSagemcomFastData = hass.data[DOMAIN][config_entry.entry_id] async_add_entities( - SagemcomScannerEntity(coordinator, idx, config_entry.entry_id) - for idx, device in coordinator.data.items() + SagemcomScannerEntity(data.coordinator, idx, config_entry.entry_id) + for idx, device in data.coordinator.data.items() )