diff --git a/homeassistant/components/airvisual_pro/__init__.py b/homeassistant/components/airvisual_pro/__init__.py index 88f05d28145..a02e735a5d6 100644 --- a/homeassistant/components/airvisual_pro/__init__.py +++ b/homeassistant/components/airvisual_pro/__init__.py @@ -38,6 +38,8 @@ PLATFORMS = [Platform.SENSOR] UPDATE_INTERVAL = timedelta(minutes=1) +AirVisualProConfigEntry = ConfigEntry["AirVisualProData"] + @dataclass class AirVisualProData: @@ -47,7 +49,9 @@ class AirVisualProData: node: NodeSamba -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry( + hass: HomeAssistant, entry: AirVisualProConfigEntry +) -> bool: """Set up AirVisual Pro from a config entry.""" node = NodeSamba(entry.data[CONF_IP_ADDRESS], entry.data[CONF_PASSWORD]) @@ -89,9 +93,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] = AirVisualProData( - coordinator=coordinator, node=node - ) + entry.runtime_data = AirVisualProData(coordinator=coordinator, node=node) async def async_shutdown(_: Event) -> None: """Define an event handler to disconnect from the websocket.""" @@ -110,11 +112,12 @@ 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: AirVisualProConfigEntry +) -> bool: """Unload a config entry.""" if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS): - data = hass.data[DOMAIN].pop(entry.entry_id) - await data.node.async_disconnect() + await entry.runtime_data.node.async_disconnect() return unload_ok diff --git a/homeassistant/components/airvisual_pro/diagnostics.py b/homeassistant/components/airvisual_pro/diagnostics.py index 9fea6e59c1d..da871442547 100644 --- a/homeassistant/components/airvisual_pro/diagnostics.py +++ b/homeassistant/components/airvisual_pro/diagnostics.py @@ -5,12 +5,10 @@ from __future__ import annotations from typing import Any from homeassistant.components.diagnostics import async_redact_data -from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_PASSWORD from homeassistant.core import HomeAssistant -from . import AirVisualProData -from .const import DOMAIN +from . import AirVisualProConfigEntry CONF_MAC_ADDRESS = "mac_address" CONF_SERIAL_NUMBER = "serial_number" @@ -23,15 +21,13 @@ TO_REDACT = { async def async_get_config_entry_diagnostics( - hass: HomeAssistant, entry: ConfigEntry + hass: HomeAssistant, entry: AirVisualProConfigEntry ) -> dict[str, Any]: """Return diagnostics for a config entry.""" - data: AirVisualProData = hass.data[DOMAIN][entry.entry_id] - return async_redact_data( { "entry": entry.as_dict(), - "data": data.coordinator.data, + "data": entry.runtime_data.coordinator.data, }, TO_REDACT, ) diff --git a/homeassistant/components/airvisual_pro/sensor.py b/homeassistant/components/airvisual_pro/sensor.py index d53def57959..895ba7d3244 100644 --- a/homeassistant/components/airvisual_pro/sensor.py +++ b/homeassistant/components/airvisual_pro/sensor.py @@ -12,7 +12,6 @@ from homeassistant.components.sensor import ( SensorEntityDescription, SensorStateClass, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( CONCENTRATION_MICROGRAMS_PER_CUBIC_METER, CONCENTRATION_PARTS_PER_MILLION, @@ -23,8 +22,7 @@ from homeassistant.const import ( from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddEntitiesCallback -from . import AirVisualProData, AirVisualProEntity -from .const import DOMAIN +from . import AirVisualProConfigEntry, AirVisualProEntity @dataclass(frozen=True, kw_only=True) @@ -129,13 +127,13 @@ def async_get_aqi_locale(settings: dict[str, Any]) -> str: async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: AirVisualProConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up AirVisual sensors based on a config entry.""" - data: AirVisualProData = hass.data[DOMAIN][entry.entry_id] - async_add_entities( - AirVisualProSensor(data.coordinator, description) + AirVisualProSensor(entry.runtime_data.coordinator, description) for description in SENSOR_DESCRIPTIONS ) diff --git a/tests/components/airvisual_pro/conftest.py b/tests/components/airvisual_pro/conftest.py index c90eb432c25..164264634b8 100644 --- a/tests/components/airvisual_pro/conftest.py +++ b/tests/components/airvisual_pro/conftest.py @@ -81,7 +81,7 @@ async def setup_airvisual_pro_fixture(hass, config, pro): return_value=pro, ), patch("homeassistant.components.airvisual_pro.NodeSamba", return_value=pro), - patch("homeassistant.components.airvisual.PLATFORMS", []), + patch("homeassistant.components.airvisual_pro.PLATFORMS", []), ): assert await async_setup_component(hass, DOMAIN, config) await hass.async_block_till_done()