Stop coordinator before connection in nibe_heatpump (#80396)

Stop coordinator in nibe_heatpump
This commit is contained in:
Joakim Plate 2022-10-17 04:06:48 +02:00 committed by GitHub
parent d537968517
commit 551fb44975
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -1,6 +1,7 @@
"""The Nibe Heat Pump integration.""" """The Nibe Heat Pump integration."""
from __future__ import annotations from __future__ import annotations
import asyncio
from collections import defaultdict from collections import defaultdict
from collections.abc import Callable, Iterable from collections.abc import Callable, Iterable
from datetime import timedelta from datetime import timedelta
@ -103,7 +104,7 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Unload a config entry.""" """Unload a config entry."""
if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS): if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
coordinator: Coordinator = hass.data[DOMAIN].pop(entry.entry_id) coordinator: Coordinator = hass.data[DOMAIN].pop(entry.entry_id)
await coordinator.connection.stop() await coordinator.async_shutdown()
return unload_ok return unload_ok
@ -173,6 +174,7 @@ class Coordinator(ContextCoordinator[dict[int, Coil], int]):
self.seed: dict[int, Coil] = {} self.seed: dict[int, Coil] = {}
self.connection = connection self.connection = connection
self.heatpump = heatpump self.heatpump = heatpump
self.task: asyncio.Task | None = None
heatpump.subscribe(heatpump.COIL_UPDATE_EVENT, self._on_coil_update) heatpump.subscribe(heatpump.COIL_UPDATE_EVENT, self._on_coil_update)
@ -219,6 +221,13 @@ class Coordinator(ContextCoordinator[dict[int, Coil], int]):
self.async_update_context_listeners([coil.address]) self.async_update_context_listeners([coil.address])
async def _async_update_data(self) -> dict[int, Coil]: async def _async_update_data(self) -> dict[int, Coil]:
self.task = asyncio.current_task()
try:
return await self._async_update_data_internal()
finally:
self.task = None
async def _async_update_data_internal(self) -> dict[int, Coil]:
@retry( @retry(
retry=retry_if_exception_type(CoilReadException), retry=retry_if_exception_type(CoilReadException),
stop=stop_after_attempt(COIL_READ_RETRIES), stop=stop_after_attempt(COIL_READ_RETRIES),
@ -249,6 +258,14 @@ class Coordinator(ContextCoordinator[dict[int, Coil], int]):
return result return result
async def async_shutdown(self):
"""Make sure a coordinator is shut down as well as it's connection."""
if self.task:
self.task.cancel()
await asyncio.wait((self.task,))
self._unschedule_refresh()
await self.connection.stop()
class CoilEntity(CoordinatorEntity[Coordinator]): class CoilEntity(CoordinatorEntity[Coordinator]):
"""Base for coil based entities.""" """Base for coil based entities."""