diff --git a/homeassistant/components/elgato/__init__.py b/homeassistant/components/elgato/__init__.py index 8d6af325213..7b331dfed66 100644 --- a/homeassistant/components/elgato/__init__.py +++ b/homeassistant/components/elgato/__init__.py @@ -4,25 +4,24 @@ from homeassistant.config_entries import ConfigEntry from homeassistant.const import Platform from homeassistant.core import HomeAssistant -from .const import DOMAIN from .coordinator import ElgatoDataUpdateCoordinator PLATFORMS = [Platform.BUTTON, Platform.LIGHT, Platform.SENSOR, Platform.SWITCH] +ElgatorConfigEntry = ConfigEntry[ElgatoDataUpdateCoordinator] -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: + +async def async_setup_entry(hass: HomeAssistant, entry: ElgatorConfigEntry) -> bool: """Set up Elgato Light from a config entry.""" coordinator = ElgatoDataUpdateCoordinator(hass, entry) await coordinator.async_config_entry_first_refresh() - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator + entry.runtime_data = coordinator await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) return True -async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_unload_entry(hass: HomeAssistant, entry: ElgatorConfigEntry) -> bool: """Unload Elgato Light config entry.""" - if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS): - del hass.data[DOMAIN][entry.entry_id] - return unload_ok + return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) diff --git a/homeassistant/components/elgato/button.py b/homeassistant/components/elgato/button.py index 47e24ca245a..aefff0b750b 100644 --- a/homeassistant/components/elgato/button.py +++ b/homeassistant/components/elgato/button.py @@ -13,13 +13,12 @@ from homeassistant.components.button import ( ButtonEntity, ButtonEntityDescription, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from . import ElgatorConfigEntry from .coordinator import ElgatoDataUpdateCoordinator from .entity import ElgatoEntity @@ -49,11 +48,11 @@ BUTTONS = [ async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: ElgatorConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up Elgato button based on a config entry.""" - coordinator: ElgatoDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data async_add_entities( ElgatoButtonEntity( coordinator=coordinator, diff --git a/homeassistant/components/elgato/diagnostics.py b/homeassistant/components/elgato/diagnostics.py index 91f5c9a8319..ac3ea0a155d 100644 --- a/homeassistant/components/elgato/diagnostics.py +++ b/homeassistant/components/elgato/diagnostics.py @@ -4,18 +4,16 @@ from __future__ import annotations from typing import Any -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant -from .const import DOMAIN -from .coordinator import ElgatoDataUpdateCoordinator +from . import ElgatorConfigEntry async def async_get_config_entry_diagnostics( - hass: HomeAssistant, entry: ConfigEntry + hass: HomeAssistant, entry: ElgatorConfigEntry ) -> dict[str, Any]: """Return diagnostics for a config entry.""" - coordinator: ElgatoDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data return { "info": coordinator.data.info.to_dict(), "state": coordinator.data.state.to_dict(), diff --git a/homeassistant/components/elgato/light.py b/homeassistant/components/elgato/light.py index 100a04fb6fb..2cd3d611bf5 100644 --- a/homeassistant/components/elgato/light.py +++ b/homeassistant/components/elgato/light.py @@ -13,7 +13,6 @@ from homeassistant.components.light import ( ColorMode, LightEntity, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.entity_platform import ( @@ -21,7 +20,8 @@ from homeassistant.helpers.entity_platform import ( async_get_current_platform, ) -from .const import DOMAIN, SERVICE_IDENTIFY +from . import ElgatorConfigEntry +from .const import SERVICE_IDENTIFY from .coordinator import ElgatoDataUpdateCoordinator from .entity import ElgatoEntity @@ -30,11 +30,11 @@ PARALLEL_UPDATES = 1 async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: ElgatorConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up Elgato Light based on a config entry.""" - coordinator: ElgatoDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data async_add_entities([ElgatoLight(coordinator)]) platform = async_get_current_platform() diff --git a/homeassistant/components/elgato/sensor.py b/homeassistant/components/elgato/sensor.py index 76d88df3fb9..f794d26cf7f 100644 --- a/homeassistant/components/elgato/sensor.py +++ b/homeassistant/components/elgato/sensor.py @@ -11,7 +11,6 @@ from homeassistant.components.sensor import ( SensorEntityDescription, SensorStateClass, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( PERCENTAGE, EntityCategory, @@ -22,7 +21,7 @@ from homeassistant.const import ( from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from . import ElgatorConfigEntry from .coordinator import ElgatoData, ElgatoDataUpdateCoordinator from .entity import ElgatoEntity @@ -102,11 +101,11 @@ SENSORS = [ async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: ElgatorConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up Elgato sensor based on a config entry.""" - coordinator: ElgatoDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data async_add_entities( ElgatoSensorEntity( diff --git a/homeassistant/components/elgato/switch.py b/homeassistant/components/elgato/switch.py index 0d20ae95e03..fe177616034 100644 --- a/homeassistant/components/elgato/switch.py +++ b/homeassistant/components/elgato/switch.py @@ -9,13 +9,12 @@ from typing import Any from elgato import Elgato, ElgatoError from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription -from homeassistant.config_entries import ConfigEntry from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from . import ElgatorConfigEntry from .coordinator import ElgatoData, ElgatoDataUpdateCoordinator from .entity import ElgatoEntity @@ -53,11 +52,11 @@ SWITCHES = [ async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: ElgatorConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up Elgato switches based on a config entry.""" - coordinator: ElgatoDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data async_add_entities( ElgatoSwitchEntity( diff --git a/tests/components/elgato/test_init.py b/tests/components/elgato/test_init.py index a4ccb302461..a6ff923beed 100644 --- a/tests/components/elgato/test_init.py +++ b/tests/components/elgato/test_init.py @@ -4,7 +4,6 @@ from unittest.mock import MagicMock from elgato import ElgatoConnectionError -from homeassistant.components.elgato.const import DOMAIN from homeassistant.config_entries import ConfigEntryState from homeassistant.core import HomeAssistant @@ -27,7 +26,6 @@ async def test_load_unload_config_entry( await hass.config_entries.async_unload(mock_config_entry.entry_id) await hass.async_block_till_done() - assert not hass.data.get(DOMAIN) assert mock_config_entry.state is ConfigEntryState.NOT_LOADED