diff --git a/homeassistant/components/energenie_power_sockets/__init__.py b/homeassistant/components/energenie_power_sockets/__init__.py index 12ddb0d1389..0496c6f9b92 100644 --- a/homeassistant/components/energenie_power_sockets/__init__.py +++ b/homeassistant/components/energenie_power_sockets/__init__.py @@ -8,12 +8,14 @@ from homeassistant.const import Platform from homeassistant.core import HomeAssistant from homeassistant.exceptions import ConfigEntryError, ConfigEntryNotReady -from .const import CONF_DEVICE_API_ID, DOMAIN +from .const import CONF_DEVICE_API_ID PLATFORMS = [Platform.SWITCH] +type EnergenieConfigEntry = ConfigEntry[PowerStripUSB] -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: + +async def async_setup_entry(hass: HomeAssistant, entry: EnergenieConfigEntry) -> bool: """Set up Energenie Power Sockets.""" try: powerstrip: PowerStripUSB | None = get_device(entry.data[CONF_DEVICE_API_ID]) @@ -26,19 +28,13 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: "Can't access Energenie Power Sockets, will retry later." ) - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = powerstrip + entry.runtime_data = powerstrip + entry.async_on_unload(powerstrip.release) 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: EnergenieConfigEntry) -> bool: """Unload config entry.""" - if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS): - powerstrip = hass.data[DOMAIN].pop(entry.entry_id) - powerstrip.release() - - if not hass.data[DOMAIN]: - hass.data.pop(DOMAIN) - - return unload_ok + return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) diff --git a/homeassistant/components/energenie_power_sockets/switch.py b/homeassistant/components/energenie_power_sockets/switch.py index 1d5b9ed5197..e4fb7653e5e 100644 --- a/homeassistant/components/energenie_power_sockets/switch.py +++ b/homeassistant/components/energenie_power_sockets/switch.py @@ -7,22 +7,22 @@ from pyegps.exceptions import EgpsException from pyegps.powerstrip import PowerStrip from homeassistant.components.switch import SwitchDeviceClass, SwitchEntity -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.device_registry import DeviceInfo from homeassistant.helpers.entity_platform import AddEntitiesCallback +from . import EnergenieConfigEntry from .const import DOMAIN async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: EnergenieConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Add EGPS sockets for passed config_entry in HA.""" - powerstrip: PowerStrip = hass.data[DOMAIN][config_entry.entry_id] + powerstrip = config_entry.runtime_data async_add_entities( ( diff --git a/tests/components/energenie_power_sockets/conftest.py b/tests/components/energenie_power_sockets/conftest.py index c142e436fd3..d0301034cf8 100644 --- a/tests/components/energenie_power_sockets/conftest.py +++ b/tests/components/energenie_power_sockets/conftest.py @@ -44,7 +44,7 @@ def get_pyegps_device_mock() -> MagicMock: fkObj = FakePowerStrip( devId=DEMO_CONFIG_DATA[CONF_DEVICE_API_ID], number_of_sockets=4 ) - fkObj.release = lambda: True + fkObj.release = lambda: None fkObj._status = [0, 1, 0, 1] usb_device_mock = MagicMock(wraps=fkObj) diff --git a/tests/components/energenie_power_sockets/test_init.py b/tests/components/energenie_power_sockets/test_init.py index 4e2fe51665b..a11cef319b2 100644 --- a/tests/components/energenie_power_sockets/test_init.py +++ b/tests/components/energenie_power_sockets/test_init.py @@ -4,7 +4,6 @@ from unittest.mock import MagicMock from pyegps.exceptions import UsbError -from homeassistant.components.energenie_power_sockets.const import DOMAIN from homeassistant.config_entries import ConfigEntryState from homeassistant.core import HomeAssistant @@ -24,13 +23,11 @@ async def test_load_unload_entry( await hass.async_block_till_done() assert entry.state is ConfigEntryState.LOADED - assert entry.entry_id in hass.data[DOMAIN] assert await hass.config_entries.async_unload(entry.entry_id) await hass.async_block_till_done() assert entry.state is ConfigEntryState.NOT_LOADED - assert DOMAIN not in hass.data async def test_device_not_found_on_load_entry( diff --git a/tests/components/energenie_power_sockets/test_switch.py b/tests/components/energenie_power_sockets/test_switch.py index 4cd2bd60028..27f13390a83 100644 --- a/tests/components/energenie_power_sockets/test_switch.py +++ b/tests/components/energenie_power_sockets/test_switch.py @@ -6,7 +6,6 @@ from pyegps.exceptions import EgpsException import pytest from syrupy.assertion import SnapshotAssertion -from homeassistant.components.energenie_power_sockets.const import DOMAIN from homeassistant.components.homeassistant import ( DOMAIN as HOME_ASSISTANT_DOMAIN, SERVICE_UPDATE_ENTITY, @@ -118,7 +117,6 @@ async def test_switch_setup( await hass.async_block_till_done() assert entry.state is ConfigEntryState.LOADED - assert entry.entry_id in hass.data[DOMAIN] state = hass.states.get(f"switch.{entity_name}") assert state == snapshot