From d15d001cfc2a55534e4cbb2fa9a81fb999f35c97 Mon Sep 17 00:00:00 2001
From: Richard Kroegel <42204099+rikroe@users.noreply.github.com>
Date: Sun, 30 Jun 2024 14:51:39 +0200
Subject: [PATCH] Use runtime_data for BMW (#120837)

---
 .../bmw_connected_drive/__init__.py           | 32 ++++++++-----------
 .../bmw_connected_drive/binary_sensor.py      |  9 +++---
 .../components/bmw_connected_drive/button.py  |  8 ++---
 .../bmw_connected_drive/device_tracker.py     |  9 +++---
 .../bmw_connected_drive/diagnostics.py        | 13 ++++----
 .../components/bmw_connected_drive/lock.py    |  8 ++---
 .../components/bmw_connected_drive/notify.py  | 13 +++++---
 .../components/bmw_connected_drive/number.py  |  8 ++---
 .../components/bmw_connected_drive/select.py  |  8 ++---
 .../components/bmw_connected_drive/sensor.py  |  8 ++---
 .../components/bmw_connected_drive/switch.py  |  8 ++---
 .../bmw_connected_drive/test_config_flow.py   |  4 +--
 .../bmw_connected_drive/test_coordinator.py   |  9 ++----
 13 files changed, 59 insertions(+), 78 deletions(-)

diff --git a/homeassistant/components/bmw_connected_drive/__init__.py b/homeassistant/components/bmw_connected_drive/__init__.py
index bd4e1cf7360..5900bd1ecc3 100644
--- a/homeassistant/components/bmw_connected_drive/__init__.py
+++ b/homeassistant/components/bmw_connected_drive/__init__.py
@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from dataclasses import dataclass
 import logging
 from typing import Any
 
@@ -18,10 +19,9 @@ from homeassistant.helpers import (
 )
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.device_registry import DeviceInfo
-from homeassistant.helpers.typing import ConfigType
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
-from .const import ATTR_VIN, ATTRIBUTION, CONF_READ_ONLY, DATA_HASS_CONFIG, DOMAIN
+from .const import ATTR_VIN, ATTRIBUTION, CONF_READ_ONLY, DOMAIN
 from .coordinator import BMWDataUpdateCoordinator
 
 _LOGGER = logging.getLogger(__name__)
@@ -55,13 +55,14 @@ PLATFORMS = [
 SERVICE_UPDATE_STATE = "update_state"
 
 
-async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
-    """Set up the BMW Connected Drive component from configuration.yaml."""
-    # Store full yaml config in data for platform.NOTIFY
-    hass.data.setdefault(DOMAIN, {})
-    hass.data[DOMAIN][DATA_HASS_CONFIG] = config
+type BMWConfigEntry = ConfigEntry[BMWData]
 
-    return True
+
+@dataclass
+class BMWData:
+    """Class to store BMW runtime data."""
+
+    coordinator: BMWDataUpdateCoordinator
 
 
 @callback
@@ -82,7 +83,7 @@ def _async_migrate_options_from_data_if_missing(
 
 
 async def _async_migrate_entries(
-    hass: HomeAssistant, config_entry: ConfigEntry
+    hass: HomeAssistant, config_entry: BMWConfigEntry
 ) -> bool:
     """Migrate old entry."""
     entity_registry = er.async_get(hass)
@@ -134,8 +135,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     )
     await coordinator.async_config_entry_first_refresh()
 
-    hass.data.setdefault(DOMAIN, {})
-    hass.data[DOMAIN][entry.entry_id] = coordinator
+    entry.runtime_data = BMWData(coordinator)
 
     # Set up all platforms except notify
     await hass.config_entries.async_forward_entry_setups(
@@ -150,7 +150,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             Platform.NOTIFY,
             DOMAIN,
             {CONF_NAME: DOMAIN, CONF_ENTITY_ID: entry.entry_id},
-            hass.data[DOMAIN][DATA_HASS_CONFIG],
+            {},
         )
     )
 
@@ -171,15 +171,11 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
-    unload_ok = await hass.config_entries.async_unload_platforms(
+
+    return await hass.config_entries.async_unload_platforms(
         entry, [platform for platform in PLATFORMS if platform != Platform.NOTIFY]
     )
 
-    if unload_ok:
-        hass.data[DOMAIN].pop(entry.entry_id)
-
-    return unload_ok
-
 
 class BMWBaseEntity(CoordinatorEntity[BMWDataUpdateCoordinator]):
     """Common base for BMW entities."""
diff --git a/homeassistant/components/bmw_connected_drive/binary_sensor.py b/homeassistant/components/bmw_connected_drive/binary_sensor.py
index d40d85e4cd4..e09241d99e7 100644
--- a/homeassistant/components/bmw_connected_drive/binary_sensor.py
+++ b/homeassistant/components/bmw_connected_drive/binary_sensor.py
@@ -17,13 +17,12 @@ from homeassistant.components.binary_sensor import (
     BinarySensorEntity,
     BinarySensorEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util.unit_system import UnitSystem
 
-from . import BMWBaseEntity
-from .const import DOMAIN, UNIT_MAP
+from . import BMWBaseEntity, BMWConfigEntry
+from .const import UNIT_MAP
 from .coordinator import BMWDataUpdateCoordinator
 
 _LOGGER = logging.getLogger(__name__)
@@ -197,11 +196,11 @@ SENSOR_TYPES: tuple[BMWBinarySensorEntityDescription, ...] = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: BMWConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the BMW binary sensors from config entry."""
-    coordinator: BMWDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator = config_entry.runtime_data.coordinator
 
     entities = [
         BMWBinarySensor(coordinator, vehicle, description, hass.config.units)
diff --git a/homeassistant/components/bmw_connected_drive/button.py b/homeassistant/components/bmw_connected_drive/button.py
index fe103f0e003..ec0212cc189 100644
--- a/homeassistant/components/bmw_connected_drive/button.py
+++ b/homeassistant/components/bmw_connected_drive/button.py
@@ -12,13 +12,11 @@ from bimmer_connected.vehicle import MyBMWVehicle
 from bimmer_connected.vehicle.remote_services import RemoteServiceStatus
 
 from homeassistant.components.button import ButtonEntity, ButtonEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import BMWBaseEntity
-from .const import DOMAIN
+from . import BMWBaseEntity, BMWConfigEntry
 
 if TYPE_CHECKING:
     from .coordinator import BMWDataUpdateCoordinator
@@ -68,11 +66,11 @@ BUTTON_TYPES: tuple[BMWButtonEntityDescription, ...] = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: BMWConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the BMW buttons from config entry."""
-    coordinator: BMWDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator = config_entry.runtime_data.coordinator
 
     entities: list[BMWButton] = []
 
diff --git a/homeassistant/components/bmw_connected_drive/device_tracker.py b/homeassistant/components/bmw_connected_drive/device_tracker.py
index d6846d0b88e..6dc54e9473f 100644
--- a/homeassistant/components/bmw_connected_drive/device_tracker.py
+++ b/homeassistant/components/bmw_connected_drive/device_tracker.py
@@ -8,12 +8,11 @@ from typing import Any
 from bimmer_connected.vehicle import MyBMWVehicle
 
 from homeassistant.components.device_tracker import SourceType, TrackerEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import BMWBaseEntity
-from .const import ATTR_DIRECTION, DOMAIN
+from . import BMWBaseEntity, BMWConfigEntry
+from .const import ATTR_DIRECTION
 from .coordinator import BMWDataUpdateCoordinator
 
 _LOGGER = logging.getLogger(__name__)
@@ -21,11 +20,11 @@ _LOGGER = logging.getLogger(__name__)
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: BMWConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the MyBMW tracker from config entry."""
-    coordinator: BMWDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator = config_entry.runtime_data.coordinator
     entities: list[BMWDeviceTracker] = []
 
     for vehicle in coordinator.account.vehicles:
diff --git a/homeassistant/components/bmw_connected_drive/diagnostics.py b/homeassistant/components/bmw_connected_drive/diagnostics.py
index c2bd4b6d24a..a3a8f5f942e 100644
--- a/homeassistant/components/bmw_connected_drive/diagnostics.py
+++ b/homeassistant/components/bmw_connected_drive/diagnostics.py
@@ -9,17 +9,16 @@ from typing import TYPE_CHECKING, Any
 from bimmer_connected.utils import MyBMWJSONEncoder
 
 from homeassistant.components.diagnostics.util import async_redact_data
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.device_registry import DeviceEntry
 
-from .const import CONF_REFRESH_TOKEN, DOMAIN
+from . import BMWConfigEntry
+from .const import CONF_REFRESH_TOKEN
 
 if TYPE_CHECKING:
     from bimmer_connected.vehicle import MyBMWVehicle
 
-    from .coordinator import BMWDataUpdateCoordinator
 
 TO_REDACT_INFO = [CONF_USERNAME, CONF_PASSWORD, CONF_REFRESH_TOKEN]
 TO_REDACT_DATA = [
@@ -47,10 +46,10 @@ def vehicle_to_dict(vehicle: MyBMWVehicle | None) -> dict:
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, config_entry: ConfigEntry
+    hass: HomeAssistant, config_entry: BMWConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    coordinator: BMWDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator = config_entry.runtime_data.coordinator
 
     coordinator.account.config.log_responses = True
     await coordinator.account.get_vehicles(force_init=True)
@@ -73,10 +72,10 @@ async def async_get_config_entry_diagnostics(
 
 
 async def async_get_device_diagnostics(
-    hass: HomeAssistant, config_entry: ConfigEntry, device: DeviceEntry
+    hass: HomeAssistant, config_entry: BMWConfigEntry, device: DeviceEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a device."""
-    coordinator: BMWDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator = config_entry.runtime_data.coordinator
 
     coordinator.account.config.log_responses = True
     await coordinator.account.get_vehicles(force_init=True)
diff --git a/homeassistant/components/bmw_connected_drive/lock.py b/homeassistant/components/bmw_connected_drive/lock.py
index e138f31ba24..4380b736811 100644
--- a/homeassistant/components/bmw_connected_drive/lock.py
+++ b/homeassistant/components/bmw_connected_drive/lock.py
@@ -10,13 +10,11 @@ from bimmer_connected.vehicle import MyBMWVehicle
 from bimmer_connected.vehicle.doors_windows import LockState
 
 from homeassistant.components.lock import LockEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import BMWBaseEntity
-from .const import DOMAIN
+from . import BMWBaseEntity, BMWConfigEntry
 from .coordinator import BMWDataUpdateCoordinator
 
 DOOR_LOCK_STATE = "door_lock_state"
@@ -25,11 +23,11 @@ _LOGGER = logging.getLogger(__name__)
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: BMWConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the MyBMW lock from config entry."""
-    coordinator: BMWDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator = config_entry.runtime_data.coordinator
 
     if not coordinator.read_only:
         async_add_entities(
diff --git a/homeassistant/components/bmw_connected_drive/notify.py b/homeassistant/components/bmw_connected_drive/notify.py
index 84bc2d8459a..8edde1a5cd7 100644
--- a/homeassistant/components/bmw_connected_drive/notify.py
+++ b/homeassistant/components/bmw_connected_drive/notify.py
@@ -24,8 +24,7 @@ from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
 
-from .const import DOMAIN
-from .coordinator import BMWDataUpdateCoordinator
+from . import BMWConfigEntry
 
 ATTR_LAT = "lat"
 ATTR_LOCATION_ATTRIBUTES = ["street", "city", "postal_code", "country"]
@@ -42,12 +41,16 @@ def get_service(
     discovery_info: DiscoveryInfoType | None = None,
 ) -> BMWNotificationService:
     """Get the BMW notification service."""
-    coordinator: BMWDataUpdateCoordinator = hass.data[DOMAIN][
+    config_entry: BMWConfigEntry | None = hass.config_entries.async_get_entry(
         (discovery_info or {})[CONF_ENTITY_ID]
-    ]
+    )
 
     targets = {}
-    if not coordinator.read_only:
+    if (
+        config_entry
+        and (coordinator := config_entry.runtime_data.coordinator)
+        and not coordinator.read_only
+    ):
         targets.update({v.name: v for v in coordinator.account.vehicles})
     return BMWNotificationService(targets)
 
diff --git a/homeassistant/components/bmw_connected_drive/number.py b/homeassistant/components/bmw_connected_drive/number.py
index defeb3f0f56..a875e1a6974 100644
--- a/homeassistant/components/bmw_connected_drive/number.py
+++ b/homeassistant/components/bmw_connected_drive/number.py
@@ -14,13 +14,11 @@ from homeassistant.components.number import (
     NumberEntityDescription,
     NumberMode,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import BMWBaseEntity
-from .const import DOMAIN
+from . import BMWBaseEntity, BMWConfigEntry
 from .coordinator import BMWDataUpdateCoordinator
 
 _LOGGER = logging.getLogger(__name__)
@@ -56,11 +54,11 @@ NUMBER_TYPES: list[BMWNumberEntityDescription] = [
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: BMWConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the MyBMW number from config entry."""
-    coordinator: BMWDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator = config_entry.runtime_data.coordinator
 
     entities: list[BMWNumber] = []
 
diff --git a/homeassistant/components/bmw_connected_drive/select.py b/homeassistant/components/bmw_connected_drive/select.py
index 2522c6bf2a6..25b816d32b9 100644
--- a/homeassistant/components/bmw_connected_drive/select.py
+++ b/homeassistant/components/bmw_connected_drive/select.py
@@ -10,14 +10,12 @@ from bimmer_connected.vehicle import MyBMWVehicle
 from bimmer_connected.vehicle.charging_profile import ChargingMode
 
 from homeassistant.components.select import SelectEntity, SelectEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import UnitOfElectricCurrent
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import BMWBaseEntity
-from .const import DOMAIN
+from . import BMWBaseEntity, BMWConfigEntry
 from .coordinator import BMWDataUpdateCoordinator
 
 _LOGGER = logging.getLogger(__name__)
@@ -63,11 +61,11 @@ SELECT_TYPES: tuple[BMWSelectEntityDescription, ...] = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: BMWConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the MyBMW lock from config entry."""
-    coordinator: BMWDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator = config_entry.runtime_data.coordinator
 
     entities: list[BMWSelect] = []
 
diff --git a/homeassistant/components/bmw_connected_drive/sensor.py b/homeassistant/components/bmw_connected_drive/sensor.py
index 1d9737c7d5f..d4ac1e3decd 100644
--- a/homeassistant/components/bmw_connected_drive/sensor.py
+++ b/homeassistant/components/bmw_connected_drive/sensor.py
@@ -18,7 +18,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     PERCENTAGE,
     STATE_UNKNOWN,
@@ -30,8 +29,7 @@ from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util import dt as dt_util
 
-from . import BMWBaseEntity
-from .const import DOMAIN
+from . import BMWBaseEntity, BMWConfigEntry
 from .coordinator import BMWDataUpdateCoordinator
 
 _LOGGER = logging.getLogger(__name__)
@@ -171,11 +169,11 @@ SENSOR_TYPES: list[BMWSensorEntityDescription] = [
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: BMWConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the MyBMW sensors from config entry."""
-    coordinator: BMWDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator = config_entry.runtime_data.coordinator
 
     entities = [
         BMWSensor(coordinator, vehicle, description)
diff --git a/homeassistant/components/bmw_connected_drive/switch.py b/homeassistant/components/bmw_connected_drive/switch.py
index 5ee31d2efd7..2c8622433ba 100644
--- a/homeassistant/components/bmw_connected_drive/switch.py
+++ b/homeassistant/components/bmw_connected_drive/switch.py
@@ -10,13 +10,11 @@ from bimmer_connected.vehicle import MyBMWVehicle
 from bimmer_connected.vehicle.fuel_and_battery import ChargingState
 
 from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import BMWBaseEntity
-from .const import DOMAIN
+from . import BMWBaseEntity, BMWConfigEntry
 from .coordinator import BMWDataUpdateCoordinator
 
 _LOGGER = logging.getLogger(__name__)
@@ -64,11 +62,11 @@ NUMBER_TYPES: list[BMWSwitchEntityDescription] = [
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: BMWConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the MyBMW switch from config entry."""
-    coordinator: BMWDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator = config_entry.runtime_data.coordinator
 
     entities: list[BMWSwitch] = []
 
diff --git a/tests/components/bmw_connected_drive/test_config_flow.py b/tests/components/bmw_connected_drive/test_config_flow.py
index 3c7f452a011..f346cd70b26 100644
--- a/tests/components/bmw_connected_drive/test_config_flow.py
+++ b/tests/components/bmw_connected_drive/test_config_flow.py
@@ -159,7 +159,7 @@ async def test_options_flow_implementation(hass: HomeAssistant) -> None:
             CONF_READ_ONLY: True,
         }
 
-        assert len(mock_setup_entry.mock_calls) == 1
+        assert len(mock_setup_entry.mock_calls) == 2
 
 
 async def test_reauth(hass: HomeAssistant) -> None:
@@ -210,4 +210,4 @@ async def test_reauth(hass: HomeAssistant) -> None:
         assert result2["reason"] == "reauth_successful"
         assert config_entry.data == FIXTURE_COMPLETE_ENTRY
 
-        assert len(mock_setup_entry.mock_calls) == 1
+        assert len(mock_setup_entry.mock_calls) == 2
diff --git a/tests/components/bmw_connected_drive/test_coordinator.py b/tests/components/bmw_connected_drive/test_coordinator.py
index 5b3f99a9414..ca629084f6c 100644
--- a/tests/components/bmw_connected_drive/test_coordinator.py
+++ b/tests/components/bmw_connected_drive/test_coordinator.py
@@ -27,10 +27,7 @@ async def test_update_success(hass: HomeAssistant) -> None:
     await hass.config_entries.async_setup(config_entry.entry_id)
     await hass.async_block_till_done()
 
-    assert (
-        hass.data[config_entry.domain][config_entry.entry_id].last_update_success
-        is True
-    )
+    assert config_entry.runtime_data.coordinator.last_update_success is True
 
 
 @pytest.mark.usefixtures("bmw_fixture")
@@ -45,7 +42,7 @@ async def test_update_failed(
     await hass.config_entries.async_setup(config_entry.entry_id)
     await hass.async_block_till_done()
 
-    coordinator = hass.data[config_entry.domain][config_entry.entry_id]
+    coordinator = config_entry.runtime_data.coordinator
 
     assert coordinator.last_update_success is True
 
@@ -74,7 +71,7 @@ async def test_update_reauth(
     await hass.config_entries.async_setup(config_entry.entry_id)
     await hass.async_block_till_done()
 
-    coordinator = hass.data[config_entry.domain][config_entry.entry_id]
+    coordinator = config_entry.runtime_data.coordinator
 
     assert coordinator.last_update_success is True