From 84cc650bb344ac429b26067d9d7fbb3553f35936 Mon Sep 17 00:00:00 2001
From: Bouwe Westerdijk <11290930+bouwew@users.noreply.github.com>
Date: Mon, 13 May 2024 09:38:06 +0200
Subject: [PATCH] Implement runtime data for Plugwise (#117172)

---
 homeassistant/components/plugwise/__init__.py | 20 ++++++++++++-------
 .../components/plugwise/binary_sensor.py      |  9 +++------
 homeassistant/components/plugwise/climate.py  |  7 ++++---
 .../components/plugwise/coordinator.py        | 12 ++++++-----
 .../components/plugwise/diagnostics.py        |  8 +++-----
 homeassistant/components/plugwise/number.py   | 10 ++++------
 homeassistant/components/plugwise/select.py   | 10 ++++------
 homeassistant/components/plugwise/sensor.py   |  7 +++----
 homeassistant/components/plugwise/switch.py   |  8 ++++----
 .../fixtures/m_adam_cooling/all_data.json     |  3 +--
 .../fixtures/m_adam_heating/all_data.json     |  3 +--
 tests/components/plugwise/test_sensor.py      |  6 +++---
 12 files changed, 50 insertions(+), 53 deletions(-)

diff --git a/homeassistant/components/plugwise/__init__.py b/homeassistant/components/plugwise/__init__.py
index 3140e518688..bce1bd81df6 100644
--- a/homeassistant/components/plugwise/__init__.py
+++ b/homeassistant/components/plugwise/__init__.py
@@ -12,16 +12,18 @@ from homeassistant.helpers import device_registry as dr, entity_registry as er
 from .const import DOMAIN, LOGGER, PLATFORMS
 from .coordinator import PlugwiseDataUpdateCoordinator
 
+PlugwiseConfigEntry = ConfigEntry[PlugwiseDataUpdateCoordinator]
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+
+async def async_setup_entry(hass: HomeAssistant, entry: PlugwiseConfigEntry) -> bool:
     """Set up Plugwise components from a config entry."""
     await er.async_migrate_entries(hass, entry.entry_id, async_migrate_entity_entry)
 
-    coordinator = PlugwiseDataUpdateCoordinator(hass, entry)
+    coordinator = PlugwiseDataUpdateCoordinator(hass)
     await coordinator.async_config_entry_first_refresh()
     migrate_sensor_entities(hass, coordinator)
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
+    entry.runtime_data = coordinator
 
     device_registry = dr.async_get(hass)
     device_registry.async_get_or_create(
@@ -38,11 +40,9 @@ 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: PlugwiseConfigEntry) -> bool:
     """Unload the Plugwise components."""
-    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)
 
 
 @callback
@@ -59,6 +59,12 @@ def async_migrate_entity_entry(entry: er.RegistryEntry) -> dict[str, Any] | None
                 "-slave_boiler_state", "-secondary_boiler_state"
             )
         }
+    if entry.domain == Platform.SENSOR and entry.unique_id.endswith(
+        "-relative_humidity"
+    ):
+        return {
+            "new_unique_id": entry.unique_id.replace("-relative_humidity", "-humidity")
+        }
     if entry.domain == Platform.SWITCH and entry.unique_id.endswith("-plug"):
         return {"new_unique_id": entry.unique_id.replace("-plug", "-relay")}
 
diff --git a/homeassistant/components/plugwise/binary_sensor.py b/homeassistant/components/plugwise/binary_sensor.py
index 01ebc736dbe..51dbb84733e 100644
--- a/homeassistant/components/plugwise/binary_sensor.py
+++ b/homeassistant/components/plugwise/binary_sensor.py
@@ -12,12 +12,11 @@ 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 DOMAIN
+from . import PlugwiseConfigEntry
 from .coordinator import PlugwiseDataUpdateCoordinator
 from .entity import PlugwiseEntity
 
@@ -78,13 +77,11 @@ BINARY_SENSORS: tuple[PlugwiseBinarySensorEntityDescription, ...] = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    entry: PlugwiseConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Smile binary_sensors from a config entry."""
-    coordinator: PlugwiseDataUpdateCoordinator = hass.data[DOMAIN][
-        config_entry.entry_id
-    ]
+    coordinator = entry.runtime_data
 
     entities: list[PlugwiseBinarySensorEntity] = []
     for device_id, device in coordinator.data.devices.items():
diff --git a/homeassistant/components/plugwise/climate.py b/homeassistant/components/plugwise/climate.py
index 7820c86a242..73151185e72 100644
--- a/homeassistant/components/plugwise/climate.py
+++ b/homeassistant/components/plugwise/climate.py
@@ -13,12 +13,12 @@ from homeassistant.components.climate import (
     HVACAction,
     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 . import PlugwiseConfigEntry
 from .const import DOMAIN, MASTER_THERMOSTATS
 from .coordinator import PlugwiseDataUpdateCoordinator
 from .entity import PlugwiseEntity
@@ -27,11 +27,12 @@ from .util import plugwise_command
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    entry: PlugwiseConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Smile Thermostats from a config entry."""
-    coordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator = entry.runtime_data
+
     async_add_entities(
         PlugwiseClimateEntity(coordinator, device_id)
         for device_id, device in coordinator.data.devices.items()
diff --git a/homeassistant/components/plugwise/coordinator.py b/homeassistant/components/plugwise/coordinator.py
index 15a0e8c4821..4cb1a35867e 100644
--- a/homeassistant/components/plugwise/coordinator.py
+++ b/homeassistant/components/plugwise/coordinator.py
@@ -27,7 +27,9 @@ class PlugwiseDataUpdateCoordinator(DataUpdateCoordinator[PlugwiseData]):
 
     _connected: bool = False
 
-    def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None:
+    config_entry: ConfigEntry
+
+    def __init__(self, hass: HomeAssistant) -> None:
         """Initialize the coordinator."""
         super().__init__(
             hass,
@@ -45,10 +47,10 @@ class PlugwiseDataUpdateCoordinator(DataUpdateCoordinator[PlugwiseData]):
         )
 
         self.api = Smile(
-            host=entry.data[CONF_HOST],
-            username=entry.data.get(CONF_USERNAME, DEFAULT_USERNAME),
-            password=entry.data[CONF_PASSWORD],
-            port=entry.data.get(CONF_PORT, DEFAULT_PORT),
+            host=self.config_entry.data[CONF_HOST],
+            username=self.config_entry.data.get(CONF_USERNAME, DEFAULT_USERNAME),
+            password=self.config_entry.data[CONF_PASSWORD],
+            port=self.config_entry.data.get(CONF_PORT, DEFAULT_PORT),
             timeout=30,
             websession=async_get_clientsession(hass, verify_ssl=False),
         )
diff --git a/homeassistant/components/plugwise/diagnostics.py b/homeassistant/components/plugwise/diagnostics.py
index 44c0fa9a1da..9d15ea4fe28 100644
--- a/homeassistant/components/plugwise/diagnostics.py
+++ b/homeassistant/components/plugwise/diagnostics.py
@@ -4,18 +4,16 @@ from __future__ import annotations
 
 from typing import Any
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 
-from .const import DOMAIN
-from .coordinator import PlugwiseDataUpdateCoordinator
+from . import PlugwiseConfigEntry
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, entry: ConfigEntry
+    hass: HomeAssistant, entry: PlugwiseConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    coordinator: PlugwiseDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator = entry.runtime_data
     return {
         "gateway": coordinator.data.gateway,
         "devices": coordinator.data.devices,
diff --git a/homeassistant/components/plugwise/number.py b/homeassistant/components/plugwise/number.py
index 2bae113a73e..ee7199cbb88 100644
--- a/homeassistant/components/plugwise/number.py
+++ b/homeassistant/components/plugwise/number.py
@@ -13,12 +13,12 @@ from homeassistant.components.number import (
     NumberEntityDescription,
     NumberMode,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory, UnitOfTemperature
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN, NumberType
+from . import PlugwiseConfigEntry
+from .const import NumberType
 from .coordinator import PlugwiseDataUpdateCoordinator
 from .entity import PlugwiseEntity
 
@@ -67,14 +67,12 @@ NUMBER_TYPES = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    entry: PlugwiseConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Plugwise number platform."""
 
-    coordinator: PlugwiseDataUpdateCoordinator = hass.data[DOMAIN][
-        config_entry.entry_id
-    ]
+    coordinator = entry.runtime_data
 
     async_add_entities(
         PlugwiseNumberEntity(coordinator, device_id, description)
diff --git a/homeassistant/components/plugwise/select.py b/homeassistant/components/plugwise/select.py
index 10718a818ff..0b370dc55d2 100644
--- a/homeassistant/components/plugwise/select.py
+++ b/homeassistant/components/plugwise/select.py
@@ -8,12 +8,12 @@ from dataclasses import dataclass
 from plugwise import Smile
 
 from homeassistant.components.select import SelectEntity, SelectEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import STATE_ON, EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN, SelectOptionsType, SelectType
+from . import PlugwiseConfigEntry
+from .const import SelectOptionsType, SelectType
 from .coordinator import PlugwiseDataUpdateCoordinator
 from .entity import PlugwiseEntity
 
@@ -60,13 +60,11 @@ SELECT_TYPES = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    entry: PlugwiseConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Smile selector from a config entry."""
-    coordinator: PlugwiseDataUpdateCoordinator = hass.data[DOMAIN][
-        config_entry.entry_id
-    ]
+    coordinator = entry.runtime_data
 
     async_add_entities(
         PlugwiseSelectEntity(coordinator, device_id, description)
diff --git a/homeassistant/components/plugwise/sensor.py b/homeassistant/components/plugwise/sensor.py
index 2dfe97a06c5..69ee52ae777 100644
--- a/homeassistant/components/plugwise/sensor.py
+++ b/homeassistant/components/plugwise/sensor.py
@@ -12,7 +12,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     LIGHT_LUX,
     PERCENTAGE,
@@ -28,7 +27,7 @@ from homeassistant.const import (
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from . import PlugwiseConfigEntry
 from .coordinator import PlugwiseDataUpdateCoordinator
 from .entity import PlugwiseEntity
 
@@ -403,11 +402,11 @@ SENSORS: tuple[PlugwiseSensorEntityDescription, ...] = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    entry: PlugwiseConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Smile sensors from a config entry."""
-    coordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator = entry.runtime_data
 
     entities: list[PlugwiseSensorEntity] = []
     for device_id, device in coordinator.data.devices.items():
diff --git a/homeassistant/components/plugwise/switch.py b/homeassistant/components/plugwise/switch.py
index 3c737e19a4a..2c4b53cfb50 100644
--- a/homeassistant/components/plugwise/switch.py
+++ b/homeassistant/components/plugwise/switch.py
@@ -12,12 +12,11 @@ from homeassistant.components.switch import (
     SwitchEntity,
     SwitchEntityDescription,
 )
-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 DOMAIN
+from . import PlugwiseConfigEntry
 from .coordinator import PlugwiseDataUpdateCoordinator
 from .entity import PlugwiseEntity
 from .util import plugwise_command
@@ -57,11 +56,12 @@ SWITCHES: tuple[PlugwiseSwitchEntityDescription, ...] = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    entry: PlugwiseConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Smile switches from a config entry."""
-    coordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator = entry.runtime_data
+
     entities: list[PlugwiseSwitchEntity] = []
     for device_id, device in coordinator.data.devices.items():
         if not (switches := device.get("switches")):
diff --git a/tests/components/plugwise/fixtures/m_adam_cooling/all_data.json b/tests/components/plugwise/fixtures/m_adam_cooling/all_data.json
index d9bf85b4701..6cd3241a637 100644
--- a/tests/components/plugwise/fixtures/m_adam_cooling/all_data.json
+++ b/tests/components/plugwise/fixtures/m_adam_cooling/all_data.json
@@ -66,8 +66,7 @@
       "model": "ThermoTouch",
       "name": "Anna",
       "preset_modes": ["no_frost", "asleep", "vacation", "home", "away"],
-      "select_schedule": "Weekschema",
-      "selected_schedule": "None",
+      "select_schedule": "None",
       "sensors": {
         "setpoint": 23.5,
         "temperature": 25.8
diff --git a/tests/components/plugwise/fixtures/m_adam_heating/all_data.json b/tests/components/plugwise/fixtures/m_adam_heating/all_data.json
index 37fc73009d3..0e9df1a5079 100644
--- a/tests/components/plugwise/fixtures/m_adam_heating/all_data.json
+++ b/tests/components/plugwise/fixtures/m_adam_heating/all_data.json
@@ -71,8 +71,7 @@
       "model": "ThermoTouch",
       "name": "Anna",
       "preset_modes": ["no_frost", "asleep", "vacation", "home", "away"],
-      "select_schedule": "Weekschema",
-      "selected_schedule": "None",
+      "select_schedule": "None",
       "sensors": {
         "setpoint": 20.0,
         "temperature": 19.1
diff --git a/tests/components/plugwise/test_sensor.py b/tests/components/plugwise/test_sensor.py
index d1df8454f4e..53de5f8c64a 100644
--- a/tests/components/plugwise/test_sensor.py
+++ b/tests/components/plugwise/test_sensor.py
@@ -3,7 +3,7 @@
 from unittest.mock import MagicMock
 
 from homeassistant.components.plugwise.const import DOMAIN
-from homeassistant.components.plugwise.sensor import DOMAIN as SENSOR_DOMAIN
+from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_component import async_update_entity
 from homeassistant.helpers.entity_registry import async_get
@@ -58,7 +58,7 @@ async def test_unique_id_migration_humidity(
     entity_registry = async_get(hass)
     # Entry to migrate
     entity_registry.async_get_or_create(
-        SENSOR_DOMAIN,
+        Platform.SENSOR,
         DOMAIN,
         "f61f1a2535f54f52ad006a3d18e459ca-relative_humidity",
         config_entry=mock_config_entry,
@@ -67,7 +67,7 @@ async def test_unique_id_migration_humidity(
     )
     # Entry not needing migration
     entity_registry.async_get_or_create(
-        SENSOR_DOMAIN,
+        Platform.SENSOR,
         DOMAIN,
         "f61f1a2535f54f52ad006a3d18e459ca-battery",
         config_entry=mock_config_entry,