diff --git a/homeassistant/components/dlink/__init__.py b/homeassistant/components/dlink/__init__.py index 80260643223..212fe2e9e21 100644 --- a/homeassistant/components/dlink/__init__.py +++ b/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) diff --git a/homeassistant/components/dlink/entity.py b/homeassistant/components/dlink/entity.py index 2a9ac0e6c12..228dfd168a5 100644 --- a/homeassistant/components/dlink/entity.py +++ b/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: diff --git a/homeassistant/components/dlink/switch.py b/homeassistant/components/dlink/switch.py index 36bfe4fb391..54322cc6875 100644 --- a/homeassistant/components/dlink/switch.py +++ b/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): diff --git a/tests/components/dlink/test_switch.py b/tests/components/dlink/test_switch.py index d070158d9fb..0460a6a918f 100644 --- a/tests/components/dlink/test_switch.py +++ b/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,