Use runtime_data in coolmaster (#136405)

* Use runtime_data in coolmaster

* Adjust test

homeassistant/components/coolmaster/__init__.py

@@ -2,18 +2,17 @@
 
 from pycoolmasternet_async import CoolMasterNet
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, CONF_PORT, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
 
-from .const import CONF_SWING_SUPPORT, DATA_COORDINATOR, DATA_INFO, DOMAIN
+from .const import CONF_SWING_SUPPORT
-from .coordinator import CoolmasterDataUpdateCoordinator
+from .coordinator import CoolmasterConfigEntry, CoolmasterDataUpdateCoordinator
 
 PLATFORMS = [Platform.BINARY_SENSOR, Platform.BUTTON, Platform.CLIMATE, Platform.SENSOR]
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: CoolmasterConfigEntry) -> bool:
     """Set up Coolmaster from a config entry."""
     host = entry.data[CONF_HOST]
     port = entry.data[CONF_PORT]
@@ -38,21 +37,14 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             raise ConfigEntryNotReady
     except OSError as error:
         raise ConfigEntryNotReady from error
-    coordinator = CoolmasterDataUpdateCoordinator(hass, coolmaster)
+    coordinator = CoolmasterDataUpdateCoordinator(hass, entry, coolmaster, info)
-    hass.data.setdefault(DOMAIN, {})
     await coordinator.async_config_entry_first_refresh()
-    hass.data[DOMAIN][entry.entry_id] = {
+    entry.runtime_data = coordinator
-        DATA_INFO: info,
-        DATA_COORDINATOR: 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: CoolmasterConfigEntry) -> bool:
     """Unload a Coolmaster config entry."""
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-    if unload_ok:
-        hass.data[DOMAIN].pop(entry.entry_id)
-    return unload_ok

homeassistant/components/coolmaster/binary_sensor.py

@@ -7,26 +7,23 @@ from homeassistant.components.binary_sensor import (
     BinarySensorEntity,
     BinarySensorEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DATA_COORDINATOR, DATA_INFO, DOMAIN
+from .coordinator import CoolmasterConfigEntry
 from .entity import CoolmasterEntity
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: CoolmasterConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the CoolMasterNet binary_sensor platform."""
-    info = hass.data[DOMAIN][config_entry.entry_id][DATA_INFO]
+    coordinator = config_entry.runtime_data
-    coordinator = hass.data[DOMAIN][config_entry.entry_id][DATA_COORDINATOR]
     async_add_entities(
-        CoolmasterCleanFilter(coordinator, unit_id, info)
+        CoolmasterCleanFilter(coordinator, unit_id) for unit_id in coordinator.data
-        for unit_id in coordinator.data
     )
 
 

homeassistant/components/coolmaster/button.py

@@ -3,26 +3,23 @@
 from __future__ import annotations
 
 from homeassistant.components.button import ButtonEntity, ButtonEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DATA_COORDINATOR, DATA_INFO, DOMAIN
+from .coordinator import CoolmasterConfigEntry
 from .entity import CoolmasterEntity
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: CoolmasterConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the CoolMasterNet button platform."""
-    info = hass.data[DOMAIN][config_entry.entry_id][DATA_INFO]
+    coordinator = config_entry.runtime_data
-    coordinator = hass.data[DOMAIN][config_entry.entry_id][DATA_COORDINATOR]
     async_add_entities(
-        CoolmasterResetFilter(coordinator, unit_id, info)
+        CoolmasterResetFilter(coordinator, unit_id) for unit_id in coordinator.data
-        for unit_id in coordinator.data
     )
 
 

homeassistant/components/coolmaster/climate.py

@@ -12,13 +12,13 @@ from homeassistant.components.climate import (
     ClimateEntityFeature,
     HVACMode,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_TEMPERATURE, UnitOfTemperature
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import CONF_SUPPORTED_MODES, DATA_COORDINATOR, DATA_INFO, DOMAIN
+from .const import CONF_SUPPORTED_MODES
+from .coordinator import CoolmasterConfigEntry, CoolmasterDataUpdateCoordinator
 from .entity import CoolmasterEntity
 
 CM_TO_HA_STATE = {
@@ -38,15 +38,16 @@ _LOGGER = logging.getLogger(__name__)
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: CoolmasterConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the CoolMasterNet climate platform."""
-    info = hass.data[DOMAIN][config_entry.entry_id][DATA_INFO]
+    coordinator = config_entry.runtime_data
-    coordinator = hass.data[DOMAIN][config_entry.entry_id][DATA_COORDINATOR]
+    supported_modes: list[str] = config_entry.data[CONF_SUPPORTED_MODES]
-    supported_modes = config_entry.data.get(CONF_SUPPORTED_MODES)
     async_add_entities(
-        CoolmasterClimate(coordinator, unit_id, info, supported_modes)
+        CoolmasterClimate(
+            coordinator, unit_id, [HVACMode(mode) for mode in supported_modes]
+        )
         for unit_id in coordinator.data
     )
 
@@ -56,9 +57,14 @@ class CoolmasterClimate(CoolmasterEntity, ClimateEntity):
 
     _attr_name = None
 
-    def __init__(self, coordinator, unit_id, info, supported_modes):
+    def __init__(
+        self,
+        coordinator: CoolmasterDataUpdateCoordinator,
+        unit_id: str,
+        supported_modes: list[HVACMode],
+    ) -> None:
         """Initialize the climate device."""
-        super().__init__(coordinator, unit_id, info)
+        super().__init__(coordinator, unit_id)
         self._attr_hvac_modes = supported_modes
         self._attr_unique_id = unit_id
 

homeassistant/components/coolmaster/const.py

@@ -1,8 +1,5 @@
 """Constants for the Coolmaster integration."""
 
-DATA_INFO = "info"
-DATA_COORDINATOR = "coordinator"
-
 DOMAIN = "coolmaster"
 
 DEFAULT_PORT = 10102

homeassistant/components/coolmaster/coordinator.py

@@ -1,8 +1,15 @@
 """DataUpdateCoordinator for coolmaster integration."""
 
+from __future__ import annotations
+
 import logging
 
+from pycoolmasternet_async import CoolMasterNet
+from pycoolmasternet_async.coolmasternet import CoolMasterNetUnit
+
 from homeassistant.components.climate import SCAN_INTERVAL
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
 from .const import DOMAIN
@@ -10,21 +17,34 @@ from .const import DOMAIN
 _LOGGER = logging.getLogger(__name__)
 
 
-class CoolmasterDataUpdateCoordinator(DataUpdateCoordinator):
+type CoolmasterConfigEntry = ConfigEntry[CoolmasterDataUpdateCoordinator]
+
+
+class CoolmasterDataUpdateCoordinator(
+    DataUpdateCoordinator[dict[str, CoolMasterNetUnit]]
+):
     """Class to manage fetching Coolmaster data."""
 
-    def __init__(self, hass, coolmaster):
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        entry: CoolmasterConfigEntry,
+        coolmaster: CoolMasterNet,
+        info: dict[str, str],
+    ) -> None:
         """Initialize global Coolmaster data updater."""
         self._coolmaster = coolmaster
+        self.info = info
 
         super().__init__(
             hass,
             _LOGGER,
+            config_entry=entry,
             name=DOMAIN,
             update_interval=SCAN_INTERVAL,
         )
 
-    async def _async_update_data(self):
+    async def _async_update_data(self) -> dict[str, CoolMasterNetUnit]:
         """Fetch data from Coolmaster."""
         try:
             return await self._coolmaster.status()

homeassistant/components/coolmaster/entity.py

@@ -1,7 +1,5 @@
 """Base entity for Coolmaster integration."""
 
-from pycoolmasternet_async.coolmasternet import CoolMasterNetUnit
-
 from homeassistant.core import callback
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
@@ -19,18 +17,17 @@ class CoolmasterEntity(CoordinatorEntity[CoolmasterDataUpdateCoordinator]):
         self,
         coordinator: CoolmasterDataUpdateCoordinator,
         unit_id: str,
-        info: dict[str, str],
     ) -> None:
         """Initiate CoolmasterEntity."""
         super().__init__(coordinator)
         self._unit_id: str = unit_id
-        self._unit: CoolMasterNetUnit = coordinator.data[self._unit_id]
+        self._unit = coordinator.data[self._unit_id]
         self._attr_device_info: DeviceInfo = DeviceInfo(
             identifiers={(DOMAIN, unit_id)},
             manufacturer="CoolAutomation",
             model="CoolMasterNet",
             name=unit_id,
-            sw_version=info["version"],
+            sw_version=coordinator.info["version"],
         )
         if hasattr(self, "entity_description"):
             self._attr_unique_id: str = f"{unit_id}-{self.entity_description.key}"

homeassistant/components/coolmaster/sensor.py

@@ -3,26 +3,23 @@
 from __future__ import annotations
 
 from homeassistant.components.sensor import SensorEntity, SensorEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DATA_COORDINATOR, DATA_INFO, DOMAIN
+from .coordinator import CoolmasterConfigEntry
 from .entity import CoolmasterEntity
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: CoolmasterConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the CoolMasterNet sensor platform."""
-    info = hass.data[DOMAIN][config_entry.entry_id][DATA_INFO]
+    coordinator = config_entry.runtime_data
-    coordinator = hass.data[DOMAIN][config_entry.entry_id][DATA_COORDINATOR]
     async_add_entities(
-        CoolmasterCleanFilter(coordinator, unit_id, info)
+        CoolmasterCleanFilter(coordinator, unit_id) for unit_id in coordinator.data
-        for unit_id in coordinator.data
     )
 
 

tests/components/coolmaster/test_init.py

@@ -1,6 +1,5 @@
 """The test for the Coolmaster integration."""
 
-from homeassistant.components.coolmaster.const import DOMAIN
 from homeassistant.config_entries import ConfigEntry, ConfigEntryState
 from homeassistant.core import HomeAssistant
 
@@ -20,8 +19,6 @@ async def test_unload_entry(
     load_int: ConfigEntry,
 ) -> None:
     """Test Coolmaster unloading an entry."""
-    assert load_int.entry_id in hass.data.get(DOMAIN)
     await hass.config_entries.async_unload(load_int.entry_id)
     await hass.async_block_till_done()
     assert load_int.state is ConfigEntryState.NOT_LOADED
-    assert not hass.data.get(DOMAIN)