From e68901235b411189e8309f17329264872afbbae7 Mon Sep 17 00:00:00 2001 From: Joost Lekkerkerker Date: Wed, 1 May 2024 22:13:38 +0200 Subject: [PATCH] Store runtime data in entry in Analytics Insights (#116441) --- .../components/analytics_insights/__init__.py | 22 +++++++++++-------- .../analytics_insights/coordinator.py | 7 ++++-- .../components/analytics_insights/sensor.py | 7 +++--- .../components/analytics_insights/conftest.py | 2 +- 4 files changed, 22 insertions(+), 16 deletions(-) diff --git a/homeassistant/components/analytics_insights/__init__.py b/homeassistant/components/analytics_insights/__init__.py index 79556fb68c2..3069e8dd12d 100644 --- a/homeassistant/components/analytics_insights/__init__.py +++ b/homeassistant/components/analytics_insights/__init__.py @@ -15,10 +15,11 @@ from homeassistant.core import HomeAssistant from homeassistant.exceptions import ConfigEntryNotReady from homeassistant.helpers.aiohttp_client import async_get_clientsession -from .const import CONF_TRACKED_INTEGRATIONS, DOMAIN +from .const import CONF_TRACKED_INTEGRATIONS from .coordinator import HomeassistantAnalyticsDataUpdateCoordinator PLATFORMS: list[Platform] = [Platform.SENSOR] +AnalyticsInsightsConfigEntry = ConfigEntry["AnalyticsInsightsData"] @dataclass(frozen=True) @@ -29,7 +30,9 @@ class AnalyticsInsightsData: names: dict[str, str] -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry( + hass: HomeAssistant, entry: AnalyticsInsightsConfigEntry +) -> bool: """Set up Homeassistant Analytics from a config entry.""" client = HomeassistantAnalyticsClient(session=async_get_clientsession(hass)) @@ -49,7 +52,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: await coordinator.async_config_entry_first_refresh() - hass.data[DOMAIN] = AnalyticsInsightsData(coordinator=coordinator, names=names) + entry.runtime_data = AnalyticsInsightsData(coordinator=coordinator, names=names) await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) entry.async_on_unload(entry.add_update_listener(update_listener)) @@ -57,14 +60,15 @@ 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: AnalyticsInsightsConfigEntry +) -> bool: """Unload a config entry.""" - if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS): - hass.data.pop(DOMAIN) - - return unload_ok + return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) -async def update_listener(hass: HomeAssistant, entry: ConfigEntry) -> None: +async def update_listener( + hass: HomeAssistant, entry: AnalyticsInsightsConfigEntry +) -> None: """Handle options update.""" await hass.config_entries.async_reload(entry.entry_id) diff --git a/homeassistant/components/analytics_insights/coordinator.py b/homeassistant/components/analytics_insights/coordinator.py index 759ce567898..2f863bf7771 100644 --- a/homeassistant/components/analytics_insights/coordinator.py +++ b/homeassistant/components/analytics_insights/coordinator.py @@ -4,6 +4,7 @@ from __future__ import annotations from dataclasses import dataclass from datetime import timedelta +from typing import TYPE_CHECKING from python_homeassistant_analytics import ( CustomIntegration, @@ -12,7 +13,6 @@ from python_homeassistant_analytics import ( HomeassistantAnalyticsNotModifiedError, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed @@ -23,6 +23,9 @@ from .const import ( LOGGER, ) +if TYPE_CHECKING: + from . import AnalyticsInsightsConfigEntry + @dataclass(frozen=True) class AnalyticsData: @@ -35,7 +38,7 @@ class AnalyticsData: class HomeassistantAnalyticsDataUpdateCoordinator(DataUpdateCoordinator[AnalyticsData]): """A Homeassistant Analytics Data Update Coordinator.""" - config_entry: ConfigEntry + config_entry: AnalyticsInsightsConfigEntry def __init__( self, hass: HomeAssistant, client: HomeassistantAnalyticsClient diff --git a/homeassistant/components/analytics_insights/sensor.py b/homeassistant/components/analytics_insights/sensor.py index ee1496eb52c..f7a77743b94 100644 --- a/homeassistant/components/analytics_insights/sensor.py +++ b/homeassistant/components/analytics_insights/sensor.py @@ -10,7 +10,6 @@ from homeassistant.components.sensor import ( SensorEntityDescription, SensorStateClass, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo @@ -18,7 +17,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.typing import StateType from homeassistant.helpers.update_coordinator import CoordinatorEntity -from . import AnalyticsInsightsData +from . import AnalyticsInsightsConfigEntry from .const import DOMAIN from .coordinator import AnalyticsData, HomeassistantAnalyticsDataUpdateCoordinator @@ -60,12 +59,12 @@ def get_custom_integration_entity_description( async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: AnalyticsInsightsConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Initialize the entries.""" - analytics_data: AnalyticsInsightsData = hass.data[DOMAIN] + analytics_data = entry.runtime_data coordinator: HomeassistantAnalyticsDataUpdateCoordinator = ( analytics_data.coordinator ) diff --git a/tests/components/analytics_insights/conftest.py b/tests/components/analytics_insights/conftest.py index 03bd24faeea..51d25f0a2cc 100644 --- a/tests/components/analytics_insights/conftest.py +++ b/tests/components/analytics_insights/conftest.py @@ -7,10 +7,10 @@ import pytest from python_homeassistant_analytics import CurrentAnalytics from python_homeassistant_analytics.models import CustomIntegration, Integration -from homeassistant.components.analytics_insights import DOMAIN from homeassistant.components.analytics_insights.const import ( CONF_TRACKED_CUSTOM_INTEGRATIONS, CONF_TRACKED_INTEGRATIONS, + DOMAIN, ) from tests.common import MockConfigEntry, load_fixture, load_json_object_fixture