Store integration data in config entry runtime_data

custom_components/hilo/__init__.py

@@ -84,6 +84,8 @@
     Platform.SWITCH,
 ]
 
+type HiloConfigEntry = ConfigEntry[Hilo]
+
 # Suffixes used when building gateway-attached entity unique_ids in sensor.py
 # (e.g. f"{device.identifier.lower()}-{suffix}"). Used to migrate those
 # unique_ids when the gateway's identifier changes from DSN-based to MAC-based.
@@ -191,7 +193,7 @@ def _async_migrate_gateway_device_identifier(
     LOG.info("Migrated gateway device identifier %s -> %s", old_dsn, new_mac)
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: HiloConfigEntry) -> bool:
     """Set up Hilo as config entry."""
     HiloFlowHandler.async_register_implementation(
         hass, AuthCodeWithPKCEImplementation(hass)
@@ -237,8 +239,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         LOG.debug("Timeout with API connection")
         raise ConfigEntryNotReady from err
 
-    hass.data.setdefault(DOMAIN, {})
-    hass.data[DOMAIN][entry.entry_id] = hilo
+    entry.runtime_data = hilo
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
@@ -308,7 +309,9 @@ async def handle_debug_event(event: Event):
         LOG.debug("HILO_DEBUG: log_traces is %s", log_traces)
         hass.bus.async_listen("hilo_debug", handle_debug_event)
 
-    async def async_reload_entry(_: HomeAssistant, updated_entry: ConfigEntry) -> None:
+    async def async_reload_entry(
+        _: HomeAssistant, updated_entry: HiloConfigEntry
+    ) -> None:
         """Handle an options update.
 
         This method will get called in two scenarios:
@@ -327,11 +330,14 @@ async def async_reload_entry(_: HomeAssistant, updated_entry: ConfigEntry) -> No
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: HiloConfigEntry) -> bool:
     """Unload a Hilo config entry."""
     LOG.debug("Unloading Hilo Integration")
 
-    hilo = hass.data[DOMAIN][entry.entry_id]
+    if not await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
+        return False
+
+    hilo = entry.runtime_data
 
     hilo.should_signalr_reconnect = False
 
@@ -360,12 +366,8 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     except Exception as err:
         LOG.error("Error disconnecting challenge SignalR hub: %s", err)
 
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-    if unload_ok:
-        LOG.debug("Hilo Integration unloaded")
-        hass.data[DOMAIN].pop(entry.entry_id)
-
-    return unload_ok
+    LOG.debug("Hilo Integration unloaded")
+    return True
 
 
 async def async_migrate_entry(hass, config_entry: ConfigEntry):

custom_components/hilo/climate.py

@@ -8,7 +8,6 @@
     HVACAction,
     HVACMode,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     ATTR_TEMPERATURE,
     PRECISION_TENTHS,
@@ -19,8 +18,8 @@
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util import slugify
 
-from . import Hilo
-from .const import CLIMATE_CLASSES, DOMAIN, LOG
+from . import Hilo, HiloConfigEntry
+from .const import CLIMATE_CLASSES, LOG
 from .entity import HiloEntity
 
 
@@ -45,10 +44,12 @@ def validate_reduction_phase(events, tag):
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: HiloConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Hilo climate entities from a config entry."""
-    hilo = hass.data[DOMAIN][entry.entry_id]
+    hilo = entry.runtime_data
     entities = []
     for d in hilo.devices.all:
         if d.type in CLIMATE_CLASSES:

custom_components/hilo/light.py

@@ -2,23 +2,24 @@
 
 from homeassistant.components.light import ATTR_BRIGHTNESS, ATTR_HS_COLOR, LightEntity
 from homeassistant.components.light.const import ColorMode
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.debounce import Debouncer
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util import slugify
 
-from . import Hilo
-from .const import DOMAIN, LIGHT_CLASSES, LOG
+from . import Hilo, HiloConfigEntry
+from .const import LIGHT_CLASSES, LOG
 from .entity import HiloEntity
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: HiloConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Hilo light entities from a config entry."""
-    hilo = hass.data[DOMAIN][entry.entry_id]
+    hilo = entry.runtime_data
     entities = []
 
     for d in hilo.devices.all:

custom_components/hilo/quality_scale.yaml

@@ -32,7 +32,7 @@ rules:
   entity-event-setup: todo
   entity-unique-id: done
   has-entity-name: todo
-  runtime-data: todo
+  runtime-data: done
   test-before-configure: done
   test-before-setup:
     status: todo
@@ -87,4 +87,4 @@ rules:
   # Platinum
   async-dependency: todo
   inject-websession: todo
-  strict-typing: todo
+  strict-typing: todo

custom_components/hilo/sensor.py

@@ -13,7 +13,6 @@
     SensorEntity,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     CONF_SCAN_INTERVAL,
     CURRENCY_DOLLAR,
@@ -54,7 +53,7 @@
 import yaml
 from yaml.scanner import ScannerError
 
-from . import Hilo
+from . import Hilo, HiloConfigEntry
 from .const import (
     CONF_ENERGY_METER_PERIOD,
     CONF_GENERATE_ENERGY_METERS,
@@ -143,10 +142,12 @@ def generate_entities_from_device(device, hilo, scan_interval):
 
 # noinspection GrazieInspection
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: HiloConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Hilo sensors based on a config entry."""
-    hilo = hass.data[DOMAIN][entry.entry_id]
+    hilo = entry.runtime_data
     new_entities = []
     cost_entities = []
     hq_plan_name = entry.options.get(CONF_HQ_PLAN_NAME, DEFAULT_HQ_PLAN_NAME)

custom_components/hilo/switch.py

@@ -1,23 +1,24 @@
 """Support for Hilo switches."""
 
 from homeassistant.components.switch import SwitchEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util import slugify
 from pyhilo.device.switch import Switch
 
-from . import Hilo
-from .const import DOMAIN, LOG, SWITCH_CLASSES
+from . import Hilo, HiloConfigEntry
+from .const import LOG, SWITCH_CLASSES
 from .entity import HiloEntity
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: HiloConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Hilo switches based on a config entry."""
-    hilo = hass.data[DOMAIN][entry.entry_id]
+    hilo = entry.runtime_data
     entities = []
 
     for d in hilo.devices.all:

tests/test_tariff.py

@@ -212,7 +212,7 @@ async def test_sensor_rates_flex_d_initialization(
     gateway.name = "Gateway"
     gateway.has_attribute.return_value = False
     hilo.devices = MagicMock(all=[gateway])
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = hilo
+    entry.runtime_data = hilo
 
     added_entities = []