From b18356bb3f78641e123ade31e1a7cf1e1cddaffa Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Thu, 1 Jun 2023 18:54:25 -0500 Subject: [PATCH] Fix august aiohttp session being closed out from under it (#93942) * Fix august aiohttp session being closed out from under it fixes #93941 * Fix august aiohttp session being closed out from under it fixes #93941 * Fix august aiohttp session being closed out from under it fixes #93941 --- homeassistant/components/august/__init__.py | 9 ++++-- .../components/august/config_flow.py | 29 +++++++++++++++++-- homeassistant/components/august/gateway.py | 10 ++----- tests/components/august/test_gateway.py | 2 +- 4 files changed, 36 insertions(+), 14 deletions(-) diff --git a/homeassistant/components/august/__init__.py b/homeassistant/components/august/__init__.py index 8be7d8dd2d1..8738b58dab9 100644 --- a/homeassistant/components/august/__init__.py +++ b/homeassistant/components/august/__init__.py @@ -23,7 +23,7 @@ from homeassistant.exceptions import ( ConfigEntryNotReady, HomeAssistantError, ) -from homeassistant.helpers import device_registry as dr, discovery_flow +from homeassistant.helpers import aiohttp_client, device_registry as dr, discovery_flow from .activity import ActivityStream from .const import CONF_BRAND, DOMAIN, MIN_TIME_BETWEEN_DETAIL_UPDATES, PLATFORMS @@ -44,8 +44,11 @@ YALEXS_BLE_DOMAIN = "yalexs_ble" async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Set up August from a config entry.""" - - august_gateway = AugustGateway(hass) + # Create an aiohttp session instead of using the default one since the + # default one is likely to trigger august's WAF if another integration + # is also using Cloudflare + session = aiohttp_client.async_create_clientsession(hass) + august_gateway = AugustGateway(hass, session) try: await august_gateway.async_setup(entry.data) diff --git a/homeassistant/components/august/config_flow.py b/homeassistant/components/august/config_flow.py index 58f1c2fc976..670d1608421 100644 --- a/homeassistant/components/august/config_flow.py +++ b/homeassistant/components/august/config_flow.py @@ -4,13 +4,16 @@ from dataclasses import dataclass import logging from typing import Any +import aiohttp import voluptuous as vol from yalexs.authenticator import ValidationResult from yalexs.const import BRANDS, DEFAULT_BRAND from homeassistant import config_entries from homeassistant.const import CONF_PASSWORD, CONF_USERNAME +from homeassistant.core import callback from homeassistant.data_entry_flow import FlowResult +from homeassistant.helpers import aiohttp_client from .const import ( CONF_ACCESS_TOKEN_CACHE_FILE, @@ -80,6 +83,7 @@ class AugustConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): def __init__(self): """Store an AugustGateway().""" self._august_gateway: AugustGateway | None = None + self._aiohttp_session: aiohttp.ClientSession | None = None self._user_auth_details: dict[str, Any] = {} self._needs_reset = True self._mode = None @@ -87,7 +91,6 @@ class AugustConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): async def async_step_user(self, user_input=None): """Handle the initial step.""" - self._august_gateway = AugustGateway(self.hass) return await self.async_step_user_validate() async def async_step_user_validate(self, user_input=None): @@ -151,12 +154,30 @@ class AugustConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): }, ) + @callback + def _async_get_gateway(self) -> AugustGateway: + """Set up the gateway.""" + if self._august_gateway is not None: + return self._august_gateway + # Create an aiohttp session instead of using the default one since the + # default one is likely to trigger august's WAF if another integration + # is also using Cloudflare + self._aiohttp_session = aiohttp_client.async_create_clientsession(self.hass) + self._august_gateway = AugustGateway(self.hass, self._aiohttp_session) + return self._august_gateway + + @callback + def _async_shutdown_gateway(self) -> None: + """Shutdown the gateway.""" + if self._aiohttp_session is not None: + self._aiohttp_session.detach() + self._august_gateway = None + async def async_step_reauth(self, entry_data: Mapping[str, Any]) -> FlowResult: """Handle configuration by re-auth.""" self._user_auth_details = dict(entry_data) self._mode = "reauth" self._needs_reset = True - self._august_gateway = AugustGateway(self.hass) return await self.async_step_reauth_validate() async def async_step_reauth_validate(self, user_input=None): @@ -206,7 +227,7 @@ class AugustConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): async def _async_auth_or_validate(self) -> ValidateResult: """Authenticate or validate.""" user_auth_details = self._user_auth_details - gateway = self._august_gateway + gateway = self._async_get_gateway() assert gateway is not None await self._async_reset_access_token_cache_if_needed( gateway, @@ -239,6 +260,8 @@ class AugustConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): async def _async_update_or_create_entry(self, info: dict[str, Any]) -> FlowResult: """Update existing entry or create a new one.""" + self._async_shutdown_gateway() + existing_entry = await self.async_set_unique_id( self._user_auth_details[CONF_USERNAME] ) diff --git a/homeassistant/components/august/gateway.py b/homeassistant/components/august/gateway.py index 9dcf96f057a..badff721d10 100644 --- a/homeassistant/components/august/gateway.py +++ b/homeassistant/components/august/gateway.py @@ -7,7 +7,7 @@ import logging import os from typing import Any -from aiohttp import ClientError, ClientResponseError +from aiohttp import ClientError, ClientResponseError, ClientSession from yalexs.api_async import ApiAsync from yalexs.authenticator_async import AuthenticationState, AuthenticatorAsync from yalexs.authenticator_common import Authentication @@ -16,7 +16,6 @@ from yalexs.exceptions import AugustApiAIOHTTPError from homeassistant.const import CONF_PASSWORD, CONF_TIMEOUT, CONF_USERNAME from homeassistant.core import HomeAssistant, callback -from homeassistant.helpers import aiohttp_client from .const import ( CONF_ACCESS_TOKEN_CACHE_FILE, @@ -35,12 +34,9 @@ _LOGGER = logging.getLogger(__name__) class AugustGateway: """Handle the connection to August.""" - def __init__(self, hass: HomeAssistant) -> None: + def __init__(self, hass: HomeAssistant, aiohttp_session: ClientSession) -> None: """Init the connection.""" - # Create an aiohttp session instead of using the default one since the - # default one is likely to trigger august's WAF if another integration - # is also using Cloudflare - self._aiohttp_session = aiohttp_client.async_create_clientsession(hass) + self._aiohttp_session = aiohttp_session self._token_refresh_lock = asyncio.Lock() self._access_token_cache_file: str | None = None self._hass: HomeAssistant = hass diff --git a/tests/components/august/test_gateway.py b/tests/components/august/test_gateway.py index d0e18c0bed4..2a364304c4b 100644 --- a/tests/components/august/test_gateway.py +++ b/tests/components/august/test_gateway.py @@ -35,7 +35,7 @@ async def _patched_refresh_access_token( "original_token", 1234, AuthenticationState.AUTHENTICATED ) ) - august_gateway = AugustGateway(hass) + august_gateway = AugustGateway(hass, MagicMock()) mocked_config = _mock_get_config() await august_gateway.async_setup(mocked_config[DOMAIN]) await august_gateway.async_authenticate()