Use runtime_data for roku (home-assistant#132781)

* use runtime_data for roku * unload cleanup * tweaks * tweaks * fix tests * fix tests * Update config_flow.py * Update config_flow.py
2Fake · Dec 11, 2024 · 73feeac · 73feeac
1 parent 355e80a
commit 73feeac
Show file tree

Hide file tree

Showing 9 changed files with 34 additions and 66 deletions.
diff --git a/homeassistant/components/roku/__init__.py b/homeassistant/components/roku/__init__.py
@@ -6,7 +6,7 @@
 from homeassistant.const import CONF_HOST, Platform
 from homeassistant.core import HomeAssistant
 
-from .const import CONF_PLAY_MEDIA_APP_ID, DEFAULT_PLAY_MEDIA_APP_ID, DOMAIN
+from .const import CONF_PLAY_MEDIA_APP_ID, DEFAULT_PLAY_MEDIA_APP_ID
 from .coordinator import RokuDataUpdateCoordinator
 
 PLATFORMS = [
@@ -17,8 +17,10 @@
     Platform.SENSOR,
 ]
 
+type RokuConfigEntry = ConfigEntry[RokuDataUpdateCoordinator]
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+
+async def async_setup_entry(hass: HomeAssistant, entry: RokuConfigEntry) -> bool:
     """Set up Roku from a config entry."""
     if (device_id := entry.unique_id) is None:
         device_id = entry.entry_id
@@ -33,7 +35,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     )
     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)
 
@@ -42,13 +44,11 @@ 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: RokuConfigEntry) -> bool:
     """Unload a config entry."""
-    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-        hass.data[DOMAIN].pop(entry.entry_id)
-    return unload_ok
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
 
 
-async def async_reload_entry(hass: HomeAssistant, entry: ConfigEntry) -> None:
+async def async_reload_entry(hass: HomeAssistant, entry: RokuConfigEntry) -> None:
     """Reload the config entry when it changed."""
     await hass.config_entries.async_reload(entry.entry_id)
diff --git a/homeassistant/components/roku/binary_sensor.py b/homeassistant/components/roku/binary_sensor.py
@@ -11,12 +11,11 @@
     BinarySensorEntity,
     BinarySensorEntityDescription,
 )
-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 RokuConfigEntry
 from .entity import RokuEntity
 
 
@@ -56,15 +55,13 @@ class RokuBinarySensorEntityDescription(BinarySensorEntityDescription):
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: RokuConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up a Roku binary sensors based on a config entry."""
-    coordinator = hass.data[DOMAIN][entry.entry_id]
-
     async_add_entities(
         RokuBinarySensorEntity(
-            coordinator=coordinator,
+            coordinator=entry.runtime_data,
             description=description,
         )
         for description in BINARY_SENSORS

diff --git a/homeassistant/components/roku/config_flow.py b/homeassistant/components/roku/config_flow.py
@@ -10,16 +10,12 @@
 import voluptuous as vol
 
 from homeassistant.components import ssdp, zeroconf
-from homeassistant.config_entries import (
-    ConfigEntry,
-    ConfigFlow,
-    ConfigFlowResult,
-    OptionsFlow,
-)
+from homeassistant.config_entries import ConfigFlow, ConfigFlowResult, OptionsFlow
 from homeassistant.const import CONF_HOST, CONF_NAME
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
+from . import RokuConfigEntry
 from .const import CONF_PLAY_MEDIA_APP_ID, DEFAULT_PLAY_MEDIA_APP_ID, DOMAIN
 
 DATA_SCHEMA = vol.Schema({vol.Required(CONF_HOST): str})
@@ -164,7 +160,7 @@ async def async_step_discovery_confirm(
     @staticmethod
     @callback
     def async_get_options_flow(
-        config_entry: ConfigEntry,
+        config_entry: RokuConfigEntry,
     ) -> RokuOptionsFlowHandler:
         """Create the options flow."""
         return RokuOptionsFlowHandler()

diff --git a/homeassistant/components/roku/diagnostics.py b/homeassistant/components/roku/diagnostics.py
@@ -4,25 +4,21 @@
 
 from typing import Any
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 
-from .const import DOMAIN
-from .coordinator import RokuDataUpdateCoordinator
+from . import RokuConfigEntry
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, config_entry: ConfigEntry
+    hass: HomeAssistant, entry: RokuConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    coordinator: RokuDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id]
-
     return {
         "entry": {
             "data": {
-                **config_entry.data,
+                **entry.data,
             },
-            "unique_id": config_entry.unique_id,
+            "unique_id": entry.unique_id,
         },
-        "data": coordinator.data.as_dict(),
+        "data": entry.runtime_data.data.as_dict(),
     }
diff --git a/homeassistant/components/roku/media_player.py b/homeassistant/components/roku/media_player.py
@@ -23,13 +23,13 @@
     async_process_play_media_url,
 )
 from homeassistant.components.stream import FORMAT_CONTENT_TYPE, HLS_PROVIDER
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_NAME
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_platform
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import VolDictType
 
+from . import RokuConfigEntry
 from .browse_media import async_browse_media
 from .const import (
     ATTR_ARTIST_NAME,
@@ -38,7 +38,6 @@
     ATTR_KEYWORD,
     ATTR_MEDIA_TYPE,
     ATTR_THUMBNAIL,
-    DOMAIN,
     SERVICE_SEARCH,
 )
 from .coordinator import RokuDataUpdateCoordinator
@@ -83,15 +82,13 @@
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant, entry: RokuConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up the Roku config entry."""
-    coordinator: RokuDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
-
     async_add_entities(
         [
             RokuMediaPlayer(
-                coordinator=coordinator,
+                coordinator=entry.runtime_data,
             )
         ],
         True,

diff --git a/homeassistant/components/roku/remote.py b/homeassistant/components/roku/remote.py
@@ -6,28 +6,24 @@
 from typing import Any
 
 from homeassistant.components.remote import ATTR_NUM_REPEATS, RemoteEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
-from .coordinator import RokuDataUpdateCoordinator
+from . import RokuConfigEntry
 from .entity import RokuEntity
 from .helpers import roku_exception_handler
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: RokuConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Load Roku remote based on a config entry."""
-    coordinator: RokuDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
-
     async_add_entities(
         [
             RokuRemote(
-                coordinator=coordinator,
+                coordinator=entry.runtime_data,
             )
         ],
         True,

diff --git a/homeassistant/components/roku/select.py b/homeassistant/components/roku/select.py
@@ -9,12 +9,10 @@
 from rokuecp.models import Device as RokuDevice
 
 from homeassistant.components.select import SelectEntity, SelectEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
-from .coordinator import RokuDataUpdateCoordinator
+from . import RokuConfigEntry
 from .entity import RokuEntity
 from .helpers import format_channel_name, roku_exception_handler
 
@@ -108,16 +106,15 @@ class RokuSelectEntityDescription(SelectEntityDescription):
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: RokuConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Roku select based on a config entry."""
-    coordinator: RokuDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
-    device: RokuDevice = coordinator.data
+    device: RokuDevice = entry.runtime_data.data
 
     entities: list[RokuSelectEntity] = [
         RokuSelectEntity(
-            coordinator=coordinator,
+            coordinator=entry.runtime_data,
             description=description,
         )
         for description in ENTITIES
@@ -126,7 +123,7 @@ async def async_setup_entry(
     if len(device.channels) > 0:
         entities.append(
             RokuSelectEntity(
-                coordinator=coordinator,
+                coordinator=entry.runtime_data,
                 description=CHANNEL_ENTITY,
             )
         )

diff --git a/homeassistant/components/roku/sensor.py b/homeassistant/components/roku/sensor.py
@@ -8,13 +8,11 @@
 from rokuecp.models import Device as RokuDevice
 
 from homeassistant.components.sensor import SensorEntity, SensorEntityDescription
-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 .coordinator import RokuDataUpdateCoordinator
+from . import RokuConfigEntry
 from .entity import RokuEntity
 
 
@@ -43,15 +41,13 @@ class RokuSensorEntityDescription(SensorEntityDescription):
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: RokuConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Roku sensor based on a config entry."""
-    coordinator: RokuDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
-
     async_add_entities(
         RokuSensorEntity(
-            coordinator=coordinator,
+            coordinator=entry.runtime_data,
             description=description,
         )
         for description in SENSORS

diff --git a/tests/components/roku/test_init.py b/tests/components/roku/test_init.py
@@ -4,7 +4,6 @@
 
 from rokuecp import RokuConnectionError
 
-from homeassistant.components.roku.const import DOMAIN
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import HomeAssistant
 
@@ -38,12 +37,7 @@ async def test_config_entry_no_unique_id(
     await hass.config_entries.async_setup(mock_config_entry.entry_id)
     await hass.async_block_till_done()
 
-    assert mock_config_entry.entry_id in hass.data[DOMAIN]
     assert mock_config_entry.state is ConfigEntryState.LOADED
-    assert (
-        hass.data[DOMAIN][mock_config_entry.entry_id].device_id
-        == mock_config_entry.entry_id
-    )
 
 
 async def test_load_unload_config_entry(
@@ -56,10 +50,9 @@ async def test_load_unload_config_entry(
     await hass.config_entries.async_setup(mock_config_entry.entry_id)
     await hass.async_block_till_done()
 
-    assert mock_config_entry.entry_id in hass.data[DOMAIN]
     assert mock_config_entry.state is ConfigEntryState.LOADED
 
     await hass.config_entries.async_unload(mock_config_entry.entry_id)
     await hass.async_block_till_done()
-    assert mock_config_entry.entry_id not in hass.data[DOMAIN]
+
     assert mock_config_entry.state is ConfigEntryState.NOT_LOADED