Use runtime_data for airvisual_pro (#116607)

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(
+    entry.runtime_data = AirVisualProData(coordinator=coordinator, node=node)
-        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 entry.runtime_data.node.async_disconnect()
-        await data.node.async_disconnect()
 
     return unload_ok
 

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 . import AirVisualProConfigEntry
-from .const import DOMAIN
 
 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,
     )

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 . import AirVisualProConfigEntry, AirVisualProEntity
-from .const import DOMAIN
 
 
 @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
     )
 

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()