Use runtime_data in energenie_power_sockets (#136801)

* Use runtime_data in energenie_power_sockets

* Fix  tests
This commit is contained in:
epenet 2025-01-29 09:36:53 +01:00 committed by GitHub
parent 609eb00a26
commit 11671e1875
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 12 additions and 21 deletions

View File

@ -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)

View File

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

View File

@ -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)

View File

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

View File

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