diff --git a/homeassistant/components/youtube/__init__.py b/homeassistant/components/youtube/__init__.py index 32863f5a77260d..cb3a9bf03706d0 100644 --- a/homeassistant/components/youtube/__init__.py +++ b/homeassistant/components/youtube/__init__.py @@ -1,10 +1,7 @@ """Support for YouTube.""" -from __future__ import annotations - from aiohttp.client_exceptions import ClientError, ClientResponseError -from homeassistant.config_entries import ConfigEntry from homeassistant.const import Platform from homeassistant.core import HomeAssistant from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady @@ -16,13 +13,13 @@ ) from .api import AsyncConfigEntryAuth -from .const import AUTH, COORDINATOR, DOMAIN -from .coordinator import YouTubeDataUpdateCoordinator +from .const import DOMAIN +from .coordinator import YouTubeConfigEntry, YouTubeDataUpdateCoordinator PLATFORMS = [Platform.SENSOR] -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry(hass: HomeAssistant, entry: YouTubeConfigEntry) -> bool: """Set up YouTube from a config entry.""" try: implementation = await async_get_config_entry_implementation(hass, entry) @@ -49,25 +46,21 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: await delete_devices(hass, entry, coordinator) - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = { - COORDINATOR: coordinator, - AUTH: auth, - } + 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: YouTubeConfigEntry) -> 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 delete_devices( - hass: HomeAssistant, entry: ConfigEntry, coordinator: YouTubeDataUpdateCoordinator + hass: HomeAssistant, + entry: YouTubeConfigEntry, + coordinator: YouTubeDataUpdateCoordinator, ) -> None: """Delete all devices created by integration.""" channel_ids = list(coordinator.data) diff --git a/homeassistant/components/youtube/coordinator.py b/homeassistant/components/youtube/coordinator.py index 476e5bb402262f..232eb381f189e7 100644 --- a/homeassistant/components/youtube/coordinator.py +++ b/homeassistant/components/youtube/coordinator.py @@ -1,7 +1,5 @@ """DataUpdateCoordinator for the YouTube integration.""" -from __future__ import annotations - from datetime import timedelta from typing import Any @@ -14,7 +12,7 @@ from homeassistant.exceptions import ConfigEntryAuthFailed from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed -from . import AsyncConfigEntryAuth +from .api import AsyncConfigEntryAuth from .const import ( ATTR_DESCRIPTION, ATTR_LATEST_VIDEO, @@ -29,14 +27,19 @@ LOGGER, ) +type YouTubeConfigEntry = ConfigEntry[YouTubeDataUpdateCoordinator] + class YouTubeDataUpdateCoordinator(DataUpdateCoordinator[dict[str, Any]]): """A YouTube Data Update Coordinator.""" - config_entry: ConfigEntry + config_entry: YouTubeConfigEntry def __init__( - self, hass: HomeAssistant, config_entry: ConfigEntry, auth: AsyncConfigEntryAuth + self, + hass: HomeAssistant, + config_entry: YouTubeConfigEntry, + auth: AsyncConfigEntryAuth, ) -> None: """Initialize the YouTube data coordinator.""" self._auth = auth diff --git a/homeassistant/components/youtube/diagnostics.py b/homeassistant/components/youtube/diagnostics.py index 9a898b7e2de7a7..f068934105360a 100644 --- a/homeassistant/components/youtube/diagnostics.py +++ b/homeassistant/components/youtube/diagnostics.py @@ -1,25 +1,26 @@ """Diagnostics support for YouTube.""" -from __future__ import annotations - from typing import Any -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant -from .const import ATTR_DESCRIPTION, ATTR_LATEST_VIDEO, COORDINATOR, DOMAIN -from .coordinator import YouTubeDataUpdateCoordinator +from .const import ATTR_DESCRIPTION, ATTR_LATEST_VIDEO +from .coordinator import YouTubeConfigEntry async def async_get_config_entry_diagnostics( - hass: HomeAssistant, entry: ConfigEntry + hass: HomeAssistant, entry: YouTubeConfigEntry ) -> dict[str, Any]: """Return diagnostics for a config entry.""" - coordinator: YouTubeDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id][ - COORDINATOR - ] + coordinator = entry.runtime_data sensor_data: dict[str, Any] = {} for channel_id, channel_data in coordinator.data.items(): - channel_data.get(ATTR_LATEST_VIDEO, {}).pop(ATTR_DESCRIPTION) - sensor_data[channel_id] = channel_data + channel_copy = dict(channel_data) + if latest_video := channel_copy.get(ATTR_LATEST_VIDEO): + channel_copy[ATTR_LATEST_VIDEO] = { + key: value + for key, value in latest_video.items() + if key != ATTR_DESCRIPTION + } + sensor_data[channel_id] = channel_copy return sensor_data diff --git a/tests/components/youtube/test_diagnostics.py b/tests/components/youtube/test_diagnostics.py index 99d8b9d5185e75..c5cb82189a6a9e 100644 --- a/tests/components/youtube/test_diagnostics.py +++ b/tests/components/youtube/test_diagnostics.py @@ -20,5 +20,4 @@ async def test_diagnostics( """Test diagnostics.""" await setup_integration() entry = hass.config_entries.async_entries(DOMAIN)[0] - assert await get_diagnostics_for_config_entry(hass, hass_client, entry) == snapshot