Cleanup Shelly entry unload (#119748)

* Cleanup Shelly entry unload

* store platforms on runtime_data
This commit is contained in:
Shay Levy 2024-06-16 16:25:23 +03:00 committed by GitHub
parent 3a672642ea
commit b20160a465
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 36 additions and 43 deletions

View File

@ -112,8 +112,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ShellyConfigEntry) -> bo
)
return False
entry.runtime_data = ShellyEntryData()
if get_device_entry_gen(entry) in RPC_GENERATIONS:
return await _async_setup_rpc_entry(hass, entry)
@ -150,7 +148,7 @@ async def _async_setup_block_entry(
device_entry = None
sleep_period = entry.data.get(CONF_SLEEP_PERIOD)
shelly_entry_data = entry.runtime_data
runtime_data = entry.runtime_data = ShellyEntryData(BLOCK_SLEEPING_PLATFORMS)
# Some old firmware have a wrong sleep period hardcoded value.
# Following code block will force the right value for affected devices
@ -171,6 +169,7 @@ async def _async_setup_block_entry(
if sleep_period == 0:
# Not a sleeping device, finish setup
LOGGER.debug("Setting up online block device %s", entry.title)
runtime_data.platforms = PLATFORMS
try:
await device.initialize()
if not device.firmware_supported:
@ -181,24 +180,26 @@ async def _async_setup_block_entry(
except InvalidAuthError as err:
raise ConfigEntryAuthFailed(repr(err)) from err
shelly_entry_data.block = ShellyBlockCoordinator(hass, entry, device)
shelly_entry_data.block.async_setup()
shelly_entry_data.rest = ShellyRestCoordinator(hass, device, entry)
await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
runtime_data.block = ShellyBlockCoordinator(hass, entry, device)
runtime_data.block.async_setup()
runtime_data.rest = ShellyRestCoordinator(hass, device, entry)
await hass.config_entries.async_forward_entry_setups(
entry, runtime_data.platforms
)
elif sleep_period is None or device_entry is None:
# Need to get sleep info or first time sleeping device setup, wait for device
LOGGER.debug(
"Setup for device %s will resume when device is online", entry.title
)
shelly_entry_data.block = ShellyBlockCoordinator(hass, entry, device)
shelly_entry_data.block.async_setup(BLOCK_SLEEPING_PLATFORMS)
runtime_data.block = ShellyBlockCoordinator(hass, entry, device)
runtime_data.block.async_setup(runtime_data.platforms)
else:
# Restore sensors for sleeping device
LOGGER.debug("Setting up offline block device %s", entry.title)
shelly_entry_data.block = ShellyBlockCoordinator(hass, entry, device)
shelly_entry_data.block.async_setup()
runtime_data.block = ShellyBlockCoordinator(hass, entry, device)
runtime_data.block.async_setup()
await hass.config_entries.async_forward_entry_setups(
entry, BLOCK_SLEEPING_PLATFORMS
entry, runtime_data.platforms
)
ir.async_delete_issue(
@ -236,11 +237,12 @@ async def _async_setup_rpc_entry(hass: HomeAssistant, entry: ShellyConfigEntry)
device_entry = None
sleep_period = entry.data.get(CONF_SLEEP_PERIOD)
shelly_entry_data = entry.runtime_data
runtime_data = entry.runtime_data = ShellyEntryData(RPC_SLEEPING_PLATFORMS)
if sleep_period == 0:
# Not a sleeping device, finish setup
LOGGER.debug("Setting up online RPC device %s", entry.title)
runtime_data.platforms = PLATFORMS
try:
await device.initialize()
if not device.firmware_supported:
@ -251,24 +253,26 @@ async def _async_setup_rpc_entry(hass: HomeAssistant, entry: ShellyConfigEntry)
except InvalidAuthError as err:
raise ConfigEntryAuthFailed(repr(err)) from err
shelly_entry_data.rpc = ShellyRpcCoordinator(hass, entry, device)
shelly_entry_data.rpc.async_setup()
shelly_entry_data.rpc_poll = ShellyRpcPollingCoordinator(hass, entry, device)
await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
runtime_data.rpc = ShellyRpcCoordinator(hass, entry, device)
runtime_data.rpc.async_setup()
runtime_data.rpc_poll = ShellyRpcPollingCoordinator(hass, entry, device)
await hass.config_entries.async_forward_entry_setups(
entry, runtime_data.platforms
)
elif sleep_period is None or device_entry is None:
# Need to get sleep info or first time sleeping device setup, wait for device
LOGGER.debug(
"Setup for device %s will resume when device is online", entry.title
)
shelly_entry_data.rpc = ShellyRpcCoordinator(hass, entry, device)
shelly_entry_data.rpc.async_setup(RPC_SLEEPING_PLATFORMS)
runtime_data.rpc = ShellyRpcCoordinator(hass, entry, device)
runtime_data.rpc.async_setup(runtime_data.platforms)
else:
# Restore sensors for sleeping device
LOGGER.debug("Setting up offline RPC device %s", entry.title)
shelly_entry_data.rpc = ShellyRpcCoordinator(hass, entry, device)
shelly_entry_data.rpc.async_setup()
runtime_data.rpc = ShellyRpcCoordinator(hass, entry, device)
runtime_data.rpc.async_setup()
await hass.config_entries.async_forward_entry_setups(
entry, RPC_SLEEPING_PLATFORMS
entry, runtime_data.platforms
)
ir.async_delete_issue(
@ -279,21 +283,6 @@ async def _async_setup_rpc_entry(hass: HomeAssistant, entry: ShellyConfigEntry)
async def async_unload_entry(hass: HomeAssistant, entry: ShellyConfigEntry) -> bool:
"""Unload a config entry."""
shelly_entry_data = entry.runtime_data
platforms = PLATFORMS
if get_device_entry_gen(entry) in RPC_GENERATIONS:
if entry.data.get(CONF_SLEEP_PERIOD):
platforms = RPC_SLEEPING_PLATFORMS
if unload_ok := await hass.config_entries.async_unload_platforms(
entry, platforms
):
if shelly_entry_data.rpc:
await shelly_entry_data.rpc.shutdown()
return unload_ok
# delete push update issue if it exists
LOGGER.debug(
"Deleting issue %s", PUSH_UPDATE_ISSUE_ID.format(unique=entry.unique_id)
@ -302,11 +291,14 @@ async def async_unload_entry(hass: HomeAssistant, entry: ShellyConfigEntry) -> b
hass, DOMAIN, PUSH_UPDATE_ISSUE_ID.format(unique=entry.unique_id)
)
if entry.data.get(CONF_SLEEP_PERIOD):
platforms = BLOCK_SLEEPING_PLATFORMS
runtime_data = entry.runtime_data
if unload_ok := await hass.config_entries.async_unload_platforms(entry, platforms):
if shelly_entry_data.block:
await shelly_entry_data.block.shutdown()
if runtime_data.rpc:
await runtime_data.rpc.shutdown()
return unload_ok
if runtime_data.block:
await runtime_data.block.shutdown()
return await hass.config_entries.async_unload_platforms(
entry, runtime_data.platforms
)

View File

@ -72,6 +72,7 @@ from .utils import (
class ShellyEntryData:
"""Class for sharing data within a given config entry."""
platforms: list[Platform]
block: ShellyBlockCoordinator | None = None
rest: ShellyRestCoordinator | None = None
rpc: ShellyRpcCoordinator | None = None