From a89853da9df9b2716aeef3b26520981fe6decbbd Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Thu, 25 Jul 2024 05:18:24 -0500
Subject: [PATCH] Migrate switchbot to use entry.runtime_data (#122530)

---
 homeassistant/components/switchbot/__init__.py  | 17 ++++-------------
 .../components/switchbot/binary_sensor.py       | 10 +++++-----
 .../components/switchbot/coordinator.py         |  3 +++
 homeassistant/components/switchbot/cover.py     | 10 +++++-----
 .../components/switchbot/humidifier.py          | 14 +++++---------
 homeassistant/components/switchbot/light.py     |  9 +++------
 homeassistant/components/switchbot/lock.py      | 11 +++++------
 homeassistant/components/switchbot/sensor.py    | 10 +++++-----
 homeassistant/components/switchbot/switch.py    | 14 +++++---------
 9 files changed, 40 insertions(+), 58 deletions(-)

diff --git a/homeassistant/components/switchbot/__init__.py b/homeassistant/components/switchbot/__init__.py
index 7bf02ed37b6..75845d3f3ce 100644
--- a/homeassistant/components/switchbot/__init__.py
+++ b/homeassistant/components/switchbot/__init__.py
@@ -24,11 +24,10 @@ from .const import (
     CONF_RETRY_COUNT,
     CONNECTABLE_SUPPORTED_MODEL_TYPES,
     DEFAULT_RETRY_COUNT,
-    DOMAIN,
     HASS_SENSOR_TYPE_TO_SWITCHBOT_MODEL,
     SupportedModels,
 )
-from .coordinator import SwitchbotDataUpdateCoordinator
+from .coordinator import SwitchbotConfigEntry, SwitchbotDataUpdateCoordinator
 
 PLATFORMS_BY_TYPE = {
     SupportedModels.BULB.value: [Platform.SENSOR, Platform.LIGHT],
@@ -79,10 +78,9 @@ CLASS_BY_DEVICE = {
 _LOGGER = logging.getLogger(__name__)
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: SwitchbotConfigEntry) -> bool:
     """Set up Switchbot from a config entry."""
     assert entry.unique_id is not None
-    hass.data.setdefault(DOMAIN, {})
     if CONF_ADDRESS not in entry.data and CONF_MAC in entry.data:
         # Bleak uses addresses not mac addresses which are actually
         # UUIDs on some platforms (MacOS).
@@ -137,7 +135,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             retry_count=entry.options[CONF_RETRY_COUNT],
         )
 
-    coordinator = hass.data[DOMAIN][entry.entry_id] = SwitchbotDataUpdateCoordinator(
+    coordinator = entry.runtime_data = SwitchbotDataUpdateCoordinator(
         hass,
         _LOGGER,
         ble_device,
@@ -167,13 +165,6 @@ async def _async_update_listener(hass: HomeAssistant, entry: ConfigEntry) -> Non
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
     sensor_type = entry.data[CONF_SENSOR_TYPE]
-    unload_ok = await hass.config_entries.async_unload_platforms(
+    return await hass.config_entries.async_unload_platforms(
         entry, PLATFORMS_BY_TYPE[sensor_type]
     )
-
-    if unload_ok:
-        hass.data[DOMAIN].pop(entry.entry_id)
-        if not hass.config_entries.async_entries(DOMAIN):
-            hass.data.pop(DOMAIN)
-
-    return unload_ok
diff --git a/homeassistant/components/switchbot/binary_sensor.py b/homeassistant/components/switchbot/binary_sensor.py
index 92e00a65d8a..a545ffd01ce 100644
--- a/homeassistant/components/switchbot/binary_sensor.py
+++ b/homeassistant/components/switchbot/binary_sensor.py
@@ -7,13 +7,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 .coordinator import SwitchbotDataUpdateCoordinator
+from .coordinator import SwitchbotConfigEntry, SwitchbotDataUpdateCoordinator
 from .entity import SwitchbotEntity
 
 PARALLEL_UPDATES = 0
@@ -70,10 +68,12 @@ BINARY_SENSOR_TYPES: dict[str, BinarySensorEntityDescription] = {
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: SwitchbotConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Switchbot curtain based on a config entry."""
-    coordinator: SwitchbotDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator = entry.runtime_data
     async_add_entities(
         SwitchBotBinarySensor(coordinator, binary_sensor)
         for binary_sensor in coordinator.device.parsed_data
diff --git a/homeassistant/components/switchbot/coordinator.py b/homeassistant/components/switchbot/coordinator.py
index 2c68b126fa5..807132d13e8 100644
--- a/homeassistant/components/switchbot/coordinator.py
+++ b/homeassistant/components/switchbot/coordinator.py
@@ -14,6 +14,7 @@ from homeassistant.components import bluetooth
 from homeassistant.components.bluetooth.active_update_coordinator import (
     ActiveBluetoothDataUpdateCoordinator,
 )
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import CoreState, HomeAssistant, callback
 
 if TYPE_CHECKING:
@@ -24,6 +25,8 @@ _LOGGER = logging.getLogger(__name__)
 
 DEVICE_STARTUP_TIMEOUT = 30
 
+type SwitchbotConfigEntry = ConfigEntry[SwitchbotDataUpdateCoordinator]
+
 
 class SwitchbotDataUpdateCoordinator(ActiveBluetoothDataUpdateCoordinator[None]):
     """Class to manage fetching switchbot data."""
diff --git a/homeassistant/components/switchbot/cover.py b/homeassistant/components/switchbot/cover.py
index 8039ff8ec15..d2fd073cdcb 100644
--- a/homeassistant/components/switchbot/cover.py
+++ b/homeassistant/components/switchbot/cover.py
@@ -16,13 +16,11 @@ from homeassistant.components.cover import (
     CoverEntity,
     CoverEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.restore_state import RestoreEntity
 
-from .const import DOMAIN
-from .coordinator import SwitchbotDataUpdateCoordinator
+from .coordinator import SwitchbotConfigEntry, SwitchbotDataUpdateCoordinator
 from .entity import SwitchbotEntity
 
 # Initialize the logger
@@ -31,10 +29,12 @@ PARALLEL_UPDATES = 0
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: SwitchbotConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Switchbot curtain based on a config entry."""
-    coordinator: SwitchbotDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator = entry.runtime_data
     if isinstance(coordinator.device, switchbot.SwitchbotBlindTilt):
         async_add_entities([SwitchBotBlindTiltEntity(coordinator)])
     else:
diff --git a/homeassistant/components/switchbot/humidifier.py b/homeassistant/components/switchbot/humidifier.py
index 3871fcb7265..40f96577842 100644
--- a/homeassistant/components/switchbot/humidifier.py
+++ b/homeassistant/components/switchbot/humidifier.py
@@ -2,8 +2,6 @@
 
 from __future__ import annotations
 
-import logging
-
 import switchbot
 
 from homeassistant.components.humidifier import (
@@ -13,24 +11,22 @@ from homeassistant.components.humidifier import (
     HumidifierEntity,
     HumidifierEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
-from .coordinator import SwitchbotDataUpdateCoordinator
+from .coordinator import SwitchbotConfigEntry
 from .entity import SwitchbotSwitchedEntity
 
 PARALLEL_UPDATES = 0
-_LOGGER = logging.getLogger(__name__)
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: SwitchbotConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Switchbot based on a config entry."""
-    coordinator: SwitchbotDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
-    async_add_entities([SwitchBotHumidifier(coordinator)])
+    async_add_entities([SwitchBotHumidifier(entry.runtime_data)])
 
 
 class SwitchBotHumidifier(SwitchbotSwitchedEntity, HumidifierEntity):
diff --git a/homeassistant/components/switchbot/light.py b/homeassistant/components/switchbot/light.py
index 649a8b34c75..836ba1bd4f3 100644
--- a/homeassistant/components/switchbot/light.py
+++ b/homeassistant/components/switchbot/light.py
@@ -13,7 +13,6 @@ from homeassistant.components.light import (
     ColorMode,
     LightEntity,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util.color import (
@@ -21,8 +20,7 @@ from homeassistant.util.color import (
     color_temperature_mired_to_kelvin,
 )
 
-from .const import DOMAIN
-from .coordinator import SwitchbotDataUpdateCoordinator
+from .coordinator import SwitchbotConfigEntry, SwitchbotDataUpdateCoordinator
 from .entity import SwitchbotEntity
 
 SWITCHBOT_COLOR_MODE_TO_HASS = {
@@ -35,12 +33,11 @@ PARALLEL_UPDATES = 0
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: SwitchbotConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the switchbot light."""
-    coordinator: SwitchbotDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
-    async_add_entities([SwitchbotLightEntity(coordinator)])
+    async_add_entities([SwitchbotLightEntity(entry.runtime_data)])
 
 
 class SwitchbotLightEntity(SwitchbotEntity, LightEntity):
diff --git a/homeassistant/components/switchbot/lock.py b/homeassistant/components/switchbot/lock.py
index 7b58a2f5ac3..cb41d14cf66 100644
--- a/homeassistant/components/switchbot/lock.py
+++ b/homeassistant/components/switchbot/lock.py
@@ -6,21 +6,20 @@ import switchbot
 from switchbot.const import LockStatus
 
 from homeassistant.components.lock import LockEntity, LockEntityFeature
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
-from .coordinator import SwitchbotDataUpdateCoordinator
+from .coordinator import SwitchbotConfigEntry, SwitchbotDataUpdateCoordinator
 from .entity import SwitchbotEntity
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: SwitchbotConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Switchbot lock based on a config entry."""
-    coordinator: SwitchbotDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
-    async_add_entities([(SwitchBotLock(coordinator))])
+    async_add_entities([(SwitchBotLock(entry.runtime_data))])
 
 
 # noinspection PyAbstractClass
diff --git a/homeassistant/components/switchbot/sensor.py b/homeassistant/components/switchbot/sensor.py
index 2a25d84aa8d..e696f21e082 100644
--- a/homeassistant/components/switchbot/sensor.py
+++ b/homeassistant/components/switchbot/sensor.py
@@ -9,7 +9,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     PERCENTAGE,
     SIGNAL_STRENGTH_DECIBELS_MILLIWATT,
@@ -20,8 +19,7 @@ from homeassistant.const import (
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
-from .coordinator import SwitchbotDataUpdateCoordinator
+from .coordinator import SwitchbotConfigEntry, SwitchbotDataUpdateCoordinator
 from .entity import SwitchbotEntity
 
 PARALLEL_UPDATES = 0
@@ -81,10 +79,12 @@ SENSOR_TYPES: dict[str, SensorEntityDescription] = {
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: SwitchbotConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Switchbot sensor based on a config entry."""
-    coordinator: SwitchbotDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator = entry.runtime_data
     entities = [
         SwitchBotSensor(coordinator, sensor)
         for sensor in coordinator.device.parsed_data
diff --git a/homeassistant/components/switchbot/switch.py b/homeassistant/components/switchbot/switch.py
index 26ceee203aa..427496ef20c 100644
--- a/homeassistant/components/switchbot/switch.py
+++ b/homeassistant/components/switchbot/switch.py
@@ -2,33 +2,29 @@
 
 from __future__ import annotations
 
-import logging
 from typing import Any
 
 import switchbot
 
 from homeassistant.components.switch import SwitchDeviceClass, SwitchEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import STATE_ON
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.restore_state import RestoreEntity
 
-from .const import DOMAIN
-from .coordinator import SwitchbotDataUpdateCoordinator
+from .coordinator import SwitchbotConfigEntry, SwitchbotDataUpdateCoordinator
 from .entity import SwitchbotSwitchedEntity
 
-# Initialize the logger
-_LOGGER = logging.getLogger(__name__)
 PARALLEL_UPDATES = 0
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: SwitchbotConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Switchbot based on a config entry."""
-    coordinator: SwitchbotDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
-    async_add_entities([SwitchBotSwitch(coordinator)])
+    async_add_entities([SwitchBotSwitch(entry.runtime_data)])
 
 
 class SwitchBotSwitch(SwitchbotSwitchedEntity, SwitchEntity, RestoreEntity):