Use runtime_data in ws66i integration

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

homeassistant/components/ws66i/media_player.py

@@ -7,26 +7,25 @@
     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
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 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(

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]

tests/components/ws66i/test_init.py

@@ -71,12 +71,12 @@ 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)
         await hass.async_block_till_done()
 
         assert method_call.called
 
-    assert not hass.data[DOMAIN]
+    assert config_entry.state is ConfigEntryState.NOT_LOADED