diff --git a/homeassistant/components/energyzero/__init__.py b/homeassistant/components/energyzero/__init__.py index 3e1bb830cce..f7591056383 100644 --- a/homeassistant/components/energyzero/__init__.py +++ b/homeassistant/components/energyzero/__init__.py @@ -13,9 +13,11 @@ from .const import DOMAIN from .coordinator import EnergyZeroDataUpdateCoordinator from .services import async_setup_services -PLATFORMS = [Platform.SENSOR] +PLATFORMS: list[Platform] = [Platform.SENSOR] CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN) +type EnergyZeroConfigEntry = ConfigEntry[EnergyZeroDataUpdateCoordinator] + async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: """Set up EnergyZero services.""" @@ -25,7 +27,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: return True -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry(hass: HomeAssistant, entry: EnergyZeroConfigEntry) -> bool: """Set up EnergyZero from a config entry.""" coordinator = EnergyZeroDataUpdateCoordinator(hass) @@ -35,15 +37,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: await coordinator.energyzero.close() raise - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator + 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: EnergyZeroConfigEntry) -> bool: """Unload EnergyZero 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) diff --git a/homeassistant/components/energyzero/diagnostics.py b/homeassistant/components/energyzero/diagnostics.py index 35d20fee929..ee1286598e6 100644 --- a/homeassistant/components/energyzero/diagnostics.py +++ b/homeassistant/components/energyzero/diagnostics.py @@ -5,11 +5,9 @@ from __future__ import annotations from datetime import timedelta from typing import Any -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant -from . import EnergyZeroDataUpdateCoordinator -from .const import DOMAIN +from . import EnergyZeroConfigEntry from .coordinator import EnergyZeroData @@ -32,30 +30,28 @@ def get_gas_price(data: EnergyZeroData, hours: int) -> float | None: async def async_get_config_entry_diagnostics( - hass: HomeAssistant, entry: ConfigEntry + hass: HomeAssistant, entry: EnergyZeroConfigEntry ) -> dict[str, Any]: """Return diagnostics for a config entry.""" - coordinator: EnergyZeroDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] - return { "entry": { "title": entry.title, }, "energy": { - "current_hour_price": coordinator.data.energy_today.current_price, - "next_hour_price": coordinator.data.energy_today.price_at_time( - coordinator.data.energy_today.utcnow() + timedelta(hours=1) + "current_hour_price": entry.runtime_data.data.energy_today.current_price, + "next_hour_price": entry.runtime_data.data.energy_today.price_at_time( + entry.runtime_data.data.energy_today.utcnow() + timedelta(hours=1) ), - "average_price": coordinator.data.energy_today.average_price, - "max_price": coordinator.data.energy_today.extreme_prices[1], - "min_price": coordinator.data.energy_today.extreme_prices[0], - "highest_price_time": coordinator.data.energy_today.highest_price_time, - "lowest_price_time": coordinator.data.energy_today.lowest_price_time, - "percentage_of_max": coordinator.data.energy_today.pct_of_max_price, - "hours_priced_equal_or_lower": coordinator.data.energy_today.hours_priced_equal_or_lower, + "average_price": entry.runtime_data.data.energy_today.average_price, + "max_price": entry.runtime_data.data.energy_today.extreme_prices[1], + "min_price": entry.runtime_data.data.energy_today.extreme_prices[0], + "highest_price_time": entry.runtime_data.data.energy_today.highest_price_time, + "lowest_price_time": entry.runtime_data.data.energy_today.lowest_price_time, + "percentage_of_max": entry.runtime_data.data.energy_today.pct_of_max_price, + "hours_priced_equal_or_lower": entry.runtime_data.data.energy_today.hours_priced_equal_or_lower, }, "gas": { - "current_hour_price": get_gas_price(coordinator.data, 0), - "next_hour_price": get_gas_price(coordinator.data, 1), + "current_hour_price": get_gas_price(entry.runtime_data.data, 0), + "next_hour_price": get_gas_price(entry.runtime_data.data, 1), }, } diff --git a/homeassistant/components/energyzero/sensor.py b/homeassistant/components/energyzero/sensor.py index f65f7bd559c..d52da599966 100644 --- a/homeassistant/components/energyzero/sensor.py +++ b/homeassistant/components/energyzero/sensor.py @@ -13,7 +13,6 @@ from homeassistant.components.sensor import ( SensorEntityDescription, SensorStateClass, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( CURRENCY_EURO, PERCENTAGE, @@ -26,6 +25,7 @@ from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.update_coordinator import CoordinatorEntity +from . import EnergyZeroConfigEntry from .const import DOMAIN, SERVICE_TYPE_DEVICE_NAMES from .coordinator import EnergyZeroData, EnergyZeroDataUpdateCoordinator @@ -142,10 +142,12 @@ def get_gas_price(data: EnergyZeroData, hours: int) -> float | None: async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: EnergyZeroConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up EnergyZero Sensors based on a config entry.""" - coordinator: EnergyZeroDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data async_add_entities( EnergyZeroSensorEntity( coordinator=coordinator, diff --git a/homeassistant/components/energyzero/services.py b/homeassistant/components/energyzero/services.py index d98699c5c08..b281274575e 100644 --- a/homeassistant/components/energyzero/services.py +++ b/homeassistant/components/energyzero/services.py @@ -107,7 +107,7 @@ def __get_coordinator( }, ) - coordinator: EnergyZeroDataUpdateCoordinator = hass.data[DOMAIN][entry_id] + coordinator: EnergyZeroDataUpdateCoordinator = entry.runtime_data return coordinator diff --git a/tests/components/energyzero/test_init.py b/tests/components/energyzero/test_init.py index 287157026f4..f8e7e75e902 100644 --- a/tests/components/energyzero/test_init.py +++ b/tests/components/energyzero/test_init.py @@ -5,7 +5,6 @@ from unittest.mock import MagicMock, patch from energyzero import EnergyZeroConnectionError import pytest -from homeassistant.components.energyzero.const import DOMAIN from homeassistant.config_entries import ConfigEntryState from homeassistant.core import HomeAssistant @@ -26,7 +25,6 @@ async def test_load_unload_config_entry( await hass.config_entries.async_unload(mock_config_entry.entry_id) await hass.async_block_till_done() - assert not hass.data.get(DOMAIN) assert mock_config_entry.state is ConfigEntryState.NOT_LOADED