diff --git a/homeassistant/components/ws66i/__init__.py b/homeassistant/components/ws66i/__init__.py index 23a27adeb691bb..1cd62d09e5b845 100644 --- a/homeassistant/components/ws66i/__init__.py +++ b/homeassistant/components/ws66i/__init__.py @@ -6,14 +6,13 @@ from pyws66i import WS66i, get_ws66i -from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_IP_ADDRESS, EVENT_HOMEASSISTANT_STOP, Platform from homeassistant.core import HomeAssistant, callback from homeassistant.exceptions import ConfigEntryNotReady -from .const import CONF_SOURCES, DOMAIN +from .const import CONF_SOURCES from .coordinator import Ws66iDataUpdateCoordinator -from .models import SourceRep, Ws66iData +from .models import SourceRep, Ws66iConfigEntry, Ws66iData _LOGGER = logging.getLogger(__name__) @@ -56,7 +55,7 @@ def _find_zones(hass: HomeAssistant, ws66i: WS66i) -> list[int]: return zone_list -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry(hass: HomeAssistant, entry: Ws66iConfigEntry) -> bool: """Set up Soundavo WS66i 6-Zone Amplifier from a config entry.""" # Get the source names from the options flow options: dict[str, dict[str, str]] @@ -86,8 +85,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: # Fetch initial data, retry on failed poll await coordinator.async_config_entry_first_refresh() - # Create the Ws66iData data class save it to hass - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = Ws66iData( + entry.runtime_data = Ws66iData( host_ip=entry.data[CONF_IP_ADDRESS], device=ws66i, sources=source_rep, @@ -109,12 +107,10 @@ def shutdown(event): return True -async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_unload_entry(hass: HomeAssistant, entry: Ws66iConfigEntry) -> bool: """Unload a config entry.""" unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS) if unload_ok: - ws66i: WS66i = hass.data[DOMAIN][entry.entry_id].device - ws66i.close() - hass.data[DOMAIN].pop(entry.entry_id) + entry.runtime_data.device.close() return unload_ok diff --git a/homeassistant/components/ws66i/media_player.py b/homeassistant/components/ws66i/media_player.py index 36b199a1c9c064..9d62ea2f94ca92 100644 --- a/homeassistant/components/ws66i/media_player.py +++ b/homeassistant/components/ws66i/media_player.py @@ -7,7 +7,6 @@ MediaPlayerEntityFeature, MediaPlayerState, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.device_registry import DeviceInfo from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback @@ -15,18 +14,18 @@ from .const import DOMAIN, MAX_VOL from .coordinator import Ws66iDataUpdateCoordinator -from .models import Ws66iData +from .models import Ws66iConfigEntry, Ws66iData PARALLEL_UPDATES = 1 async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: Ws66iConfigEntry, async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Set up the WS66i 6-zone amplifier platform from a config entry.""" - ws66i_data: Ws66iData = hass.data[DOMAIN][config_entry.entry_id] + ws66i_data = config_entry.runtime_data # Build and add the entities from the data class async_add_entities( diff --git a/homeassistant/components/ws66i/models.py b/homeassistant/components/ws66i/models.py index 3c46d07179050c..fa9be7d639af84 100644 --- a/homeassistant/components/ws66i/models.py +++ b/homeassistant/components/ws66i/models.py @@ -6,6 +6,8 @@ from pyws66i import WS66i +from homeassistant.config_entries import ConfigEntry + from .coordinator import Ws66iDataUpdateCoordinator @@ -27,3 +29,6 @@ class Ws66iData: sources: SourceRep coordinator: Ws66iDataUpdateCoordinator zones: list[int] + + +type Ws66iConfigEntry = ConfigEntry[Ws66iData] diff --git a/tests/components/ws66i/test_init.py b/tests/components/ws66i/test_init.py index e9ec78b54daf60..08b8d3f240701b 100644 --- a/tests/components/ws66i/test_init.py +++ b/tests/components/ws66i/test_init.py @@ -71,7 +71,7 @@ async def test_unload_config_entry(hass: HomeAssistant) -> None: await hass.config_entries.async_setup(config_entry.entry_id) await hass.async_block_till_done() - assert hass.data[DOMAIN][config_entry.entry_id] + assert config_entry.state is ConfigEntryState.LOADED with patch.object(MockWs66i, "close") as method_call: await hass.config_entries.async_unload(config_entry.entry_id) @@ -79,4 +79,4 @@ async def test_unload_config_entry(hass: HomeAssistant) -> None: assert method_call.called - assert not hass.data[DOMAIN] + assert config_entry.state is ConfigEntryState.NOT_LOADED