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

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

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]
         )

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

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()