Cleanup rainbird aiohttp client cleanup on EVENT_HOMEASSISTANT_CLOSE (#112292)

* Add rainbird aiohttp client cleanup for EVENT_HOMEASSISTANT_CLOSE

* Apply suggestions from code review

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

---------

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>
This commit is contained in:
Allen Porter 2024-03-05 00:35:00 -08:00 committed by GitHub
parent 1feeeb8608
commit 2e65952eae
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

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