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
This commit is contained in:
J. Nick Koston 2023-06-01 18:54:25 -05:00 committed by GitHub
parent ba66a39668
commit b18356bb3f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 36 additions and 14 deletions

View File

@ -23,7 +23,7 @@ from homeassistant.exceptions import (
ConfigEntryNotReady, ConfigEntryNotReady,
HomeAssistantError, 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 .activity import ActivityStream
from .const import CONF_BRAND, DOMAIN, MIN_TIME_BETWEEN_DETAIL_UPDATES, PLATFORMS 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: async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Set up August from a config entry.""" """Set up August from a config entry."""
# Create an aiohttp session instead of using the default one since the
august_gateway = AugustGateway(hass) # 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: try:
await august_gateway.async_setup(entry.data) await august_gateway.async_setup(entry.data)

View File

@ -4,13 +4,16 @@ from dataclasses import dataclass
import logging import logging
from typing import Any from typing import Any
import aiohttp
import voluptuous as vol import voluptuous as vol
from yalexs.authenticator import ValidationResult from yalexs.authenticator import ValidationResult
from yalexs.const import BRANDS, DEFAULT_BRAND from yalexs.const import BRANDS, DEFAULT_BRAND
from homeassistant import config_entries from homeassistant import config_entries
from homeassistant.const import CONF_PASSWORD, CONF_USERNAME from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
from homeassistant.core import callback
from homeassistant.data_entry_flow import FlowResult from homeassistant.data_entry_flow import FlowResult
from homeassistant.helpers import aiohttp_client
from .const import ( from .const import (
CONF_ACCESS_TOKEN_CACHE_FILE, CONF_ACCESS_TOKEN_CACHE_FILE,
@ -80,6 +83,7 @@ class AugustConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
def __init__(self): def __init__(self):
"""Store an AugustGateway().""" """Store an AugustGateway()."""
self._august_gateway: AugustGateway | None = None self._august_gateway: AugustGateway | None = None
self._aiohttp_session: aiohttp.ClientSession | None = None
self._user_auth_details: dict[str, Any] = {} self._user_auth_details: dict[str, Any] = {}
self._needs_reset = True self._needs_reset = True
self._mode = None self._mode = None
@ -87,7 +91,6 @@ class AugustConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
async def async_step_user(self, user_input=None): async def async_step_user(self, user_input=None):
"""Handle the initial step.""" """Handle the initial step."""
self._august_gateway = AugustGateway(self.hass)
return await self.async_step_user_validate() return await self.async_step_user_validate()
async def async_step_user_validate(self, user_input=None): 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: async def async_step_reauth(self, entry_data: Mapping[str, Any]) -> FlowResult:
"""Handle configuration by re-auth.""" """Handle configuration by re-auth."""
self._user_auth_details = dict(entry_data) self._user_auth_details = dict(entry_data)
self._mode = "reauth" self._mode = "reauth"
self._needs_reset = True self._needs_reset = True
self._august_gateway = AugustGateway(self.hass)
return await self.async_step_reauth_validate() return await self.async_step_reauth_validate()
async def async_step_reauth_validate(self, user_input=None): 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: async def _async_auth_or_validate(self) -> ValidateResult:
"""Authenticate or validate.""" """Authenticate or validate."""
user_auth_details = self._user_auth_details user_auth_details = self._user_auth_details
gateway = self._august_gateway gateway = self._async_get_gateway()
assert gateway is not None assert gateway is not None
await self._async_reset_access_token_cache_if_needed( await self._async_reset_access_token_cache_if_needed(
gateway, 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: async def _async_update_or_create_entry(self, info: dict[str, Any]) -> FlowResult:
"""Update existing entry or create a new one.""" """Update existing entry or create a new one."""
self._async_shutdown_gateway()
existing_entry = await self.async_set_unique_id( existing_entry = await self.async_set_unique_id(
self._user_auth_details[CONF_USERNAME] self._user_auth_details[CONF_USERNAME]
) )

View File

@ -7,7 +7,7 @@ import logging
import os import os
from typing import Any from typing import Any
from aiohttp import ClientError, ClientResponseError from aiohttp import ClientError, ClientResponseError, ClientSession
from yalexs.api_async import ApiAsync from yalexs.api_async import ApiAsync
from yalexs.authenticator_async import AuthenticationState, AuthenticatorAsync from yalexs.authenticator_async import AuthenticationState, AuthenticatorAsync
from yalexs.authenticator_common import Authentication 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.const import CONF_PASSWORD, CONF_TIMEOUT, CONF_USERNAME
from homeassistant.core import HomeAssistant, callback from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers import aiohttp_client
from .const import ( from .const import (
CONF_ACCESS_TOKEN_CACHE_FILE, CONF_ACCESS_TOKEN_CACHE_FILE,
@ -35,12 +34,9 @@ _LOGGER = logging.getLogger(__name__)
class AugustGateway: class AugustGateway:
"""Handle the connection to August.""" """Handle the connection to August."""
def __init__(self, hass: HomeAssistant) -> None: def __init__(self, hass: HomeAssistant, aiohttp_session: ClientSession) -> None:
"""Init the connection.""" """Init the connection."""
# Create an aiohttp session instead of using the default one since the self._aiohttp_session = aiohttp_session
# 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._token_refresh_lock = asyncio.Lock() self._token_refresh_lock = asyncio.Lock()
self._access_token_cache_file: str | None = None self._access_token_cache_file: str | None = None
self._hass: HomeAssistant = hass self._hass: HomeAssistant = hass

View File

@ -35,7 +35,7 @@ async def _patched_refresh_access_token(
"original_token", 1234, AuthenticationState.AUTHENTICATED "original_token", 1234, AuthenticationState.AUTHENTICATED
) )
) )
august_gateway = AugustGateway(hass) august_gateway = AugustGateway(hass, MagicMock())
mocked_config = _mock_get_config() mocked_config = _mock_get_config()
await august_gateway.async_setup(mocked_config[DOMAIN]) await august_gateway.async_setup(mocked_config[DOMAIN])
await august_gateway.async_authenticate() await august_gateway.async_authenticate()