Use runtime_data in energenie_power_sockets (#136801)

* Use runtime_data in energenie_power_sockets

* Fix  tests

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):
+    return 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

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(
         (

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)

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(

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