Store runtime data inside the config entry in Dlink (#119442)

homeassistant/components/dlink/__init__.py

@@ -9,13 +9,15 @@ from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_USERNAME, Platfor
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
 
-from .const import CONF_USE_LEGACY_PROTOCOL, DOMAIN
+from .const import CONF_USE_LEGACY_PROTOCOL
 from .data import SmartPlugData
 
+type DLinkConfigEntry = ConfigEntry[SmartPlugData]
+
 PLATFORMS = [Platform.SWITCH]
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: DLinkConfigEntry) -> bool:
     """Set up D-Link Power Plug from a config entry."""
     smartplug = await hass.async_add_executor_job(
         SmartPlug,
@@ -27,14 +29,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     if not smartplug.authenticated and smartplug.use_legacy_protocol:
         raise ConfigEntryNotReady("Cannot connect/authenticate")
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = SmartPlugData(smartplug)
+    entry.runtime_data = SmartPlugData(smartplug)
     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: DLinkConfigEntry) -> 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/dlink/entity.py

@@ -2,14 +2,13 @@
 
 from __future__ import annotations
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_CONNECTIONS
 from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity import Entity, EntityDescription
 
+from . import DLinkConfigEntry
 from .const import ATTRIBUTION, DOMAIN, MANUFACTURER
-from .data import SmartPlugData
 
 
 class DLinkEntity(Entity):
@@ -20,18 +19,17 @@ class DLinkEntity(Entity):
 
     def __init__(
         self,
-        config_entry: ConfigEntry,
-        data: SmartPlugData,
+        config_entry: DLinkConfigEntry,
         description: EntityDescription,
     ) -> None:
         """Initialize a D-Link Power Plug entity."""
-        self.data = data
+        self.data = config_entry.runtime_data
         self.entity_description = description
         self._attr_unique_id = f"{config_entry.entry_id}_{description.key}"
         self._attr_device_info = DeviceInfo(
             identifiers={(DOMAIN, config_entry.entry_id)},
             manufacturer=MANUFACTURER,
-            model=data.smartplug.model_name,
+            model=self.data.smartplug.model_name,
             name=config_entry.title,
         )
         if config_entry.unique_id:

homeassistant/components/dlink/switch.py

@@ -6,12 +6,12 @@ from datetime import timedelta
 from typing import Any
 
 from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_TEMPERATURE, UnitOfTemperature
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import ATTR_TOTAL_CONSUMPTION, DOMAIN
+from . import DLinkConfigEntry
+from .const import ATTR_TOTAL_CONSUMPTION
 from .entity import DLinkEntity
 
 SCAN_INTERVAL = timedelta(minutes=2)
@@ -22,13 +22,12 @@ SWITCH_TYPE = SwitchEntityDescription(
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: DLinkConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the D-Link Power Plug switch."""
-    async_add_entities(
-        [SmartPlugSwitch(entry, hass.data[DOMAIN][entry.entry_id], SWITCH_TYPE)],
-        True,
-    )
+    async_add_entities([SmartPlugSwitch(entry, SWITCH_TYPE)], True)
 
 
 class SmartPlugSwitch(DLinkEntity, SwitchEntity):

tests/components/dlink/test_switch.py

@@ -2,7 +2,7 @@
 
 from unittest.mock import patch
 
-from homeassistant.components.dlink import DOMAIN
+from homeassistant.components.dlink.const import DOMAIN
 from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
 from homeassistant.const import (
     ATTR_ENTITY_ID,