Refactor zwave_js disconnect client helper (#51718)

This commit is contained in:
Martin Hjelmare 2021-06-12 11:35:33 +02:00 committed by GitHub
parent 4afe7de07f
commit be137b085b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -388,7 +388,7 @@ async def async_setup_entry( # noqa: C901
async def handle_ha_shutdown(event: Event) -> None: async def handle_ha_shutdown(event: Event) -> None:
"""Handle HA shutdown.""" """Handle HA shutdown."""
await disconnect_client(hass, entry, client, listen_task, platform_task) await disconnect_client(hass, entry)
listen_task = asyncio.create_task( listen_task = asyncio.create_task(
client_listen(hass, entry, client, driver_ready) client_listen(hass, entry, client, driver_ready)
@ -484,14 +484,12 @@ async def client_listen(
hass.async_create_task(hass.config_entries.async_reload(entry.entry_id)) hass.async_create_task(hass.config_entries.async_reload(entry.entry_id))
async def disconnect_client( async def disconnect_client(hass: HomeAssistant, entry: ConfigEntry) -> None:
hass: HomeAssistant,
entry: ConfigEntry,
client: ZwaveClient,
listen_task: asyncio.Task,
platform_task: asyncio.Task,
) -> None:
"""Disconnect client.""" """Disconnect client."""
data = hass.data[DOMAIN][entry.entry_id]
client: ZwaveClient = data[DATA_CLIENT]
listen_task: asyncio.Task = data[DATA_CLIENT_LISTEN_TASK]
platform_task: asyncio.Task = data[DATA_START_PLATFORM_TASK]
listen_task.cancel() listen_task.cancel()
platform_task.cancel() platform_task.cancel()
platform_setup_tasks = ( platform_setup_tasks = (
@ -509,7 +507,7 @@ async def disconnect_client(
async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Unload a config entry.""" """Unload a config entry."""
info = hass.data[DOMAIN].pop(entry.entry_id) info = hass.data[DOMAIN][entry.entry_id]
for unsub in info[DATA_UNSUBSCRIBE]: for unsub in info[DATA_UNSUBSCRIBE]:
unsub() unsub()
@ -527,13 +525,9 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
unload_ok = all(await asyncio.gather(*tasks)) unload_ok = all(await asyncio.gather(*tasks))
if DATA_CLIENT_LISTEN_TASK in info: if DATA_CLIENT_LISTEN_TASK in info:
await disconnect_client( await disconnect_client(hass, entry)
hass,
entry, hass.data[DOMAIN].pop(entry.entry_id)
info[DATA_CLIENT],
info[DATA_CLIENT_LISTEN_TASK],
platform_task=info[DATA_START_PLATFORM_TASK],
)
if entry.data.get(CONF_USE_ADDON) and entry.disabled_by: if entry.data.get(CONF_USE_ADDON) and entry.disabled_by:
addon_manager: AddonManager = get_addon_manager(hass) addon_manager: AddonManager = get_addon_manager(hass)