Use runtime_data in fitbit (#137631)

homeassistant/components/fitbit/__init__.py

@@ -7,7 +7,7 @@ from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
 from homeassistant.helpers import config_entry_oauth2_flow
 
 from . import api
-from .const import DOMAIN, FitbitScope
+from .const import FitbitScope
 from .coordinator import FitbitData, FitbitDeviceCoordinator
 from .exceptions import FitbitApiException, FitbitAuthException
 from .model import config_from_entry_data
@@ -15,10 +15,11 @@ from .model import config_from_entry_data
 PLATFORMS: list[Platform] = [Platform.SENSOR]
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
-    """Set up fitbit from a config entry."""
-    hass.data.setdefault(DOMAIN, {})
+type FitbitConfigEntry = ConfigEntry[FitbitData]
 
+
+async def async_setup_entry(hass: HomeAssistant, entry: FitbitConfigEntry) -> bool:
+    """Set up fitbit from a config entry."""
     implementation = (
         await config_entry_oauth2_flow.async_get_config_entry_implementation(
             hass, entry
@@ -41,18 +42,13 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         coordinator = FitbitDeviceCoordinator(hass, fitbit_api)
         await coordinator.async_config_entry_first_refresh()
 
-    hass.data[DOMAIN][entry.entry_id] = FitbitData(
-        api=fitbit_api, device_coordinator=coordinator
-    )
+    entry.runtime_data = FitbitData(api=fitbit_api, device_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: FitbitConfigEntry) -> bool:
     """Unload a config entry."""
-    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)

homeassistant/components/fitbit/sensor.py

@@ -14,7 +14,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     PERCENTAGE,
     EntityCategory,
@@ -29,9 +28,10 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.icon import icon_for_battery_level
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
+from . import FitbitConfigEntry
 from .api import FitbitApi
 from .const import ATTRIBUTION, BATTERY_LEVELS, DOMAIN, FitbitScope, FitbitUnitSystem
-from .coordinator import FitbitData, FitbitDeviceCoordinator
+from .coordinator import FitbitDeviceCoordinator
 from .exceptions import FitbitApiException, FitbitAuthException
 from .model import FitbitDevice, config_from_entry_data
 
@@ -131,7 +131,7 @@ class FitbitSensorEntityDescription(SensorEntityDescription):
 
 
 def _build_device_info(
-    config_entry: ConfigEntry, entity_description: FitbitSensorEntityDescription
+    config_entry: FitbitConfigEntry, entity_description: FitbitSensorEntityDescription
 ) -> DeviceInfo:
     """Build device info for sensor entities info across devices."""
     unique_id = cast(str, config_entry.unique_id)
@@ -524,12 +524,12 @@ FITBIT_RESOURCE_BATTERY_LEVEL = FitbitSensorEntityDescription(
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: FitbitConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Fitbit sensor platform."""
 
-    data: FitbitData = hass.data[DOMAIN][entry.entry_id]
+    data = entry.runtime_data
     api = data.api
 
     # These are run serially to reuse the cached user profile, not gathered
@@ -601,7 +601,7 @@ class FitbitSensor(SensorEntity):
 
     def __init__(
         self,
-        config_entry: ConfigEntry,
+        config_entry: FitbitConfigEntry,
         api: FitbitApi,
         user_profile_id: str,
         description: FitbitSensorEntityDescription,