diff --git a/homeassistant/components/rainbird/__init__.py b/homeassistant/components/rainbird/__init__.py index 2a660435e17..945c06943d3 100644 --- a/homeassistant/components/rainbird/__init__.py +++ b/homeassistant/components/rainbird/__init__.py @@ -2,12 +2,20 @@ from __future__ import annotations import logging +from typing import Any +import aiohttp from pyrainbird.async_client import AsyncRainbirdClient, AsyncRainbirdController from pyrainbird.exceptions import RainbirdApiException from homeassistant.config_entries import ConfigEntry -from homeassistant.const import CONF_HOST, CONF_MAC, CONF_PASSWORD, Platform +from homeassistant.const import ( + CONF_HOST, + CONF_MAC, + CONF_PASSWORD, + EVENT_HOMEASSISTANT_CLOSE, + Platform, +) from homeassistant.core import HomeAssistant from homeassistant.exceptions import ConfigEntryNotReady from homeassistant.helpers import device_registry as dr, entity_registry as er @@ -30,13 +38,30 @@ PLATFORMS = [ DOMAIN = "rainbird" +def _async_register_clientsession_shutdown( + hass: HomeAssistant, entry: ConfigEntry, clientsession: aiohttp.ClientSession +) -> None: + """Register cleanup hooks for the clientsession.""" + + async def _async_close_websession(*_: Any) -> None: + """Close websession.""" + await clientsession.close() + + unsub = hass.bus.async_listen_once( + EVENT_HOMEASSISTANT_CLOSE, _async_close_websession + ) + entry.async_on_unload(unsub) + entry.async_on_unload(_async_close_websession) + + async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Set up the config entry for Rain Bird.""" hass.data.setdefault(DOMAIN, {}) clientsession = async_create_clientsession() - entry.async_on_unload(clientsession.close) + _async_register_clientsession_shutdown(hass, entry, clientsession) + controller = AsyncRainbirdController( AsyncRainbirdClient( clientsession,