diff --git a/homeassistant/components/wled/__init__.py b/homeassistant/components/wled/__init__.py index 6f5bb25b162..7da551b2bb9 100644 --- a/homeassistant/components/wled/__init__.py +++ b/homeassistant/components/wled/__init__.py @@ -6,7 +6,7 @@ from homeassistant.config_entries import ConfigEntry from homeassistant.const import Platform from homeassistant.core import HomeAssistant -from .const import DOMAIN, LOGGER +from .const import LOGGER from .coordinator import WLEDDataUpdateCoordinator PLATFORMS = ( @@ -20,8 +20,10 @@ PLATFORMS = ( Platform.UPDATE, ) +WLEDConfigEntry = ConfigEntry[WLEDDataUpdateCoordinator] -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: + +async def async_setup_entry(hass: HomeAssistant, entry: WLEDConfigEntry) -> bool: """Set up WLED from a config entry.""" coordinator = WLEDDataUpdateCoordinator(hass, entry=entry) await coordinator.async_config_entry_first_refresh() @@ -36,7 +38,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: ) return False - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator + entry.runtime_data = coordinator # Set up all platforms for this device/entry. await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) @@ -47,18 +49,16 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: return True -async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_unload_entry(hass: HomeAssistant, entry: WLEDConfigEntry) -> bool: """Unload WLED config entry.""" if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS): - coordinator: WLEDDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data # Ensure disconnected and cleanup stop sub await coordinator.wled.disconnect() if coordinator.unsub: coordinator.unsub() - del hass.data[DOMAIN][entry.entry_id] - return unload_ok diff --git a/homeassistant/components/wled/binary_sensor.py b/homeassistant/components/wled/binary_sensor.py index 260c43c8ba0..cceaadd84b2 100644 --- a/homeassistant/components/wled/binary_sensor.py +++ b/homeassistant/components/wled/binary_sensor.py @@ -6,23 +6,22 @@ from homeassistant.components.binary_sensor import ( BinarySensorDeviceClass, BinarySensorEntity, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from . import WLEDConfigEntry from .coordinator import WLEDDataUpdateCoordinator from .models import WLEDEntity async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: WLEDConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up a WLED binary sensor based on a config entry.""" - coordinator: WLEDDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data async_add_entities( [ WLEDUpdateBinarySensor(coordinator), diff --git a/homeassistant/components/wled/button.py b/homeassistant/components/wled/button.py index 7d3047c7c35..3165a0cba0a 100644 --- a/homeassistant/components/wled/button.py +++ b/homeassistant/components/wled/button.py @@ -3,12 +3,11 @@ from __future__ import annotations from homeassistant.components.button import ButtonDeviceClass, ButtonEntity -from homeassistant.config_entries import ConfigEntry from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from . import WLEDConfigEntry from .coordinator import WLEDDataUpdateCoordinator from .helpers import wled_exception_handler from .models import WLEDEntity @@ -16,11 +15,11 @@ from .models import WLEDEntity async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: WLEDConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up WLED button based on a config entry.""" - coordinator: WLEDDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data async_add_entities([WLEDRestartButton(coordinator)]) diff --git a/homeassistant/components/wled/diagnostics.py b/homeassistant/components/wled/diagnostics.py index f1eed3fc0aa..e81760e0f72 100644 --- a/homeassistant/components/wled/diagnostics.py +++ b/homeassistant/components/wled/diagnostics.py @@ -5,18 +5,16 @@ from __future__ import annotations from typing import Any from homeassistant.components.diagnostics import async_redact_data -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant -from .const import DOMAIN -from .coordinator import WLEDDataUpdateCoordinator +from . import WLEDConfigEntry async def async_get_config_entry_diagnostics( - hass: HomeAssistant, entry: ConfigEntry + hass: HomeAssistant, entry: WLEDConfigEntry ) -> dict[str, Any]: """Return diagnostics for a config entry.""" - coordinator: WLEDDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data data: dict[str, Any] = { "info": async_redact_data(coordinator.data.info.__dict__, "wifi"), diff --git a/homeassistant/components/wled/light.py b/homeassistant/components/wled/light.py index 1e31f090c70..7f118db5b06 100644 --- a/homeassistant/components/wled/light.py +++ b/homeassistant/components/wled/light.py @@ -15,11 +15,11 @@ from homeassistant.components.light import ( LightEntity, LightEntityFeature, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import ATTR_COLOR_PRIMARY, ATTR_ON, ATTR_SEGMENT_ID, DOMAIN +from . import WLEDConfigEntry +from .const import ATTR_COLOR_PRIMARY, ATTR_ON, ATTR_SEGMENT_ID from .coordinator import WLEDDataUpdateCoordinator from .helpers import wled_exception_handler from .models import WLEDEntity @@ -29,11 +29,11 @@ PARALLEL_UPDATES = 1 async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: WLEDConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up WLED light based on a config entry.""" - coordinator: WLEDDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data if coordinator.keep_main_light: async_add_entities([WLEDMainLight(coordinator=coordinator)]) diff --git a/homeassistant/components/wled/number.py b/homeassistant/components/wled/number.py index e6142c1cea6..b21de71a00c 100644 --- a/homeassistant/components/wled/number.py +++ b/homeassistant/components/wled/number.py @@ -9,12 +9,12 @@ from functools import partial from wled import Segment from homeassistant.components.number import NumberEntity, NumberEntityDescription -from homeassistant.config_entries import ConfigEntry from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import ATTR_INTENSITY, ATTR_SPEED, DOMAIN +from . import WLEDConfigEntry +from .const import ATTR_INTENSITY, ATTR_SPEED from .coordinator import WLEDDataUpdateCoordinator from .helpers import wled_exception_handler from .models import WLEDEntity @@ -24,11 +24,11 @@ PARALLEL_UPDATES = 1 async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: WLEDConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up WLED number based on a config entry.""" - coordinator: WLEDDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data update_segments = partial( async_update_segments, diff --git a/homeassistant/components/wled/select.py b/homeassistant/components/wled/select.py index 755cd5746e8..abae15059cd 100644 --- a/homeassistant/components/wled/select.py +++ b/homeassistant/components/wled/select.py @@ -7,12 +7,11 @@ from functools import partial from wled import Live, Playlist, Preset from homeassistant.components.select import SelectEntity -from homeassistant.config_entries import ConfigEntry from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from . import WLEDConfigEntry from .coordinator import WLEDDataUpdateCoordinator from .helpers import wled_exception_handler from .models import WLEDEntity @@ -22,11 +21,11 @@ PARALLEL_UPDATES = 1 async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: WLEDConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up WLED select based on a config entry.""" - coordinator: WLEDDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data async_add_entities( [ diff --git a/homeassistant/components/wled/sensor.py b/homeassistant/components/wled/sensor.py index daf5748021f..aa897d6d1b9 100644 --- a/homeassistant/components/wled/sensor.py +++ b/homeassistant/components/wled/sensor.py @@ -14,7 +14,6 @@ from homeassistant.components.sensor import ( SensorEntityDescription, SensorStateClass, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( PERCENTAGE, SIGNAL_STRENGTH_DECIBELS_MILLIWATT, @@ -27,7 +26,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.typing import StateType from homeassistant.util.dt import utcnow -from .const import DOMAIN +from . import WLEDConfigEntry from .coordinator import WLEDDataUpdateCoordinator from .models import WLEDEntity @@ -128,11 +127,11 @@ SENSORS: tuple[WLEDSensorEntityDescription, ...] = ( async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: WLEDConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up WLED sensor based on a config entry.""" - coordinator: WLEDDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data async_add_entities( WLEDSensorEntity(coordinator, description) for description in SENSORS diff --git a/homeassistant/components/wled/switch.py b/homeassistant/components/wled/switch.py index a5e998ec548..305303d4254 100644 --- a/homeassistant/components/wled/switch.py +++ b/homeassistant/components/wled/switch.py @@ -6,18 +6,12 @@ from functools import partial from typing import Any from homeassistant.components.switch import SwitchEntity -from homeassistant.config_entries import ConfigEntry from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import ( - ATTR_DURATION, - ATTR_FADE, - ATTR_TARGET_BRIGHTNESS, - ATTR_UDP_PORT, - DOMAIN, -) +from . import WLEDConfigEntry +from .const import ATTR_DURATION, ATTR_FADE, ATTR_TARGET_BRIGHTNESS, ATTR_UDP_PORT from .coordinator import WLEDDataUpdateCoordinator from .helpers import wled_exception_handler from .models import WLEDEntity @@ -27,11 +21,11 @@ PARALLEL_UPDATES = 1 async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: WLEDConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up WLED switch based on a config entry.""" - coordinator: WLEDDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data async_add_entities( [ diff --git a/homeassistant/components/wled/update.py b/homeassistant/components/wled/update.py index bde2986a841..5f4036cb10c 100644 --- a/homeassistant/components/wled/update.py +++ b/homeassistant/components/wled/update.py @@ -9,11 +9,10 @@ from homeassistant.components.update import ( UpdateEntity, UpdateEntityFeature, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from . import WLEDConfigEntry from .coordinator import WLEDDataUpdateCoordinator from .helpers import wled_exception_handler from .models import WLEDEntity @@ -21,11 +20,11 @@ from .models import WLEDEntity async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: WLEDConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up WLED update based on a config entry.""" - coordinator: WLEDDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data async_add_entities([WLEDUpdateEntity(coordinator)]) diff --git a/tests/components/wled/test_init.py b/tests/components/wled/test_init.py index 6f4c47ec201..f6f1da0d41e 100644 --- a/tests/components/wled/test_init.py +++ b/tests/components/wled/test_init.py @@ -67,7 +67,7 @@ async def test_setting_unique_id( hass: HomeAssistant, init_integration: MockConfigEntry ) -> None: """Test we set unique ID if not set yet.""" - assert hass.data[DOMAIN] + assert init_integration.runtime_data assert init_integration.unique_id == "aabbccddeeff"