August oauth2 exception migration (#165397)

Co-authored-by: J. Nick Koston <nick@koston.org>

homeassistant/components/august/__init__.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 from pathlib import Path
 from typing import cast
 
-from aiohttp import ClientResponseError
+from aiohttp import ClientError
 from yalexs.exceptions import AugustApiAIOHTTPError
 from yalexs.manager.exceptions import CannotConnect, InvalidAuth, RequireValidation
 from yalexs.manager.gateway import Config as YaleXSConfig
@@ -13,7 +13,12 @@ from yalexs.manager.gateway import Config as YaleXSConfig
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EVENT_HOMEASSISTANT_STOP
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
+from homeassistant.exceptions import (
+    ConfigEntryAuthFailed,
+    ConfigEntryNotReady,
+    OAuth2TokenRequestError,
+    OAuth2TokenRequestReauthError,
+)
 from homeassistant.helpers import device_registry as dr, issue_registry as ir
 from homeassistant.helpers.config_entry_oauth2_flow import (
     ImplementationUnavailableError,
@@ -45,11 +50,18 @@ async def async_setup_entry(hass: HomeAssistant, entry: AugustConfigEntry) -> bo
     august_gateway = AugustGateway(Path(hass.config.config_dir), session, oauth_session)
     try:
         await async_setup_august(hass, entry, august_gateway)
+    except OAuth2TokenRequestReauthError as err:
+        raise ConfigEntryAuthFailed from err
     except (RequireValidation, InvalidAuth) as err:
         raise ConfigEntryAuthFailed from err
     except TimeoutError as err:
         raise ConfigEntryNotReady("Timed out connecting to august api") from err
-    except (AugustApiAIOHTTPError, ClientResponseError, CannotConnect) as err:
+    except (
+        AugustApiAIOHTTPError,
+        OAuth2TokenRequestError,
+        ClientError,
+        CannotConnect,
+    ) as err:
         raise ConfigEntryNotReady from err
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
     return True

tests/components/august/test_init.py

@@ -2,7 +2,7 @@
 
 from unittest.mock import Mock, patch
 
-from aiohttp import ClientResponseError
+from aiohttp import ClientError, ClientResponseError
 import pytest
 from yalexs.const import Brand
 from yalexs.exceptions import AugustApiAIOHTTPError, InvalidAuth
@@ -18,7 +18,11 @@ from homeassistant.const import (
     STATE_ON,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import HomeAssistantError
+from homeassistant.exceptions import (
+    HomeAssistantError,
+    OAuth2TokenRequestReauthError,
+    OAuth2TokenRequestTransientError,
+)
 from homeassistant.helpers import (
     device_registry as dr,
     entity_registry as er,
@@ -304,3 +308,59 @@ async def test_oauth_implementation_not_available(hass: HomeAssistant) -> None:
         await hass.async_block_till_done()
 
     assert entry.state is ConfigEntryState.SETUP_RETRY
+
+
+async def test_oauth_token_request_reauth_error(hass: HomeAssistant) -> None:
+    """Test OAuth token request reauth error starts a reauth flow."""
+    entry = await mock_august_config_entry(hass)
+
+    with patch(
+        "homeassistant.helpers.config_entry_oauth2_flow.OAuth2Session.async_ensure_token_valid",
+        side_effect=OAuth2TokenRequestReauthError(
+            request_info=Mock(real_url="https://auth.august.com/access_token"),
+            status=401,
+            domain=DOMAIN,
+        ),
+    ):
+        await hass.config_entries.async_setup(entry.entry_id)
+        await hass.async_block_till_done()
+
+    assert entry.state is ConfigEntryState.SETUP_ERROR
+    flows = hass.config_entries.flow.async_progress()
+    assert len(flows) == 1
+    assert flows[0]["step_id"] == "pick_implementation"
+    assert flows[0]["context"]["source"] == "reauth"
+
+
+async def test_oauth_token_request_transient_error_is_retryable(
+    hass: HomeAssistant,
+) -> None:
+    """Test OAuth token transient request error marks entry for setup retry."""
+    entry = await mock_august_config_entry(hass)
+
+    with patch(
+        "homeassistant.helpers.config_entry_oauth2_flow.OAuth2Session.async_ensure_token_valid",
+        side_effect=OAuth2TokenRequestTransientError(
+            request_info=Mock(real_url="https://auth.august.com/access_token"),
+            status=500,
+            domain=DOMAIN,
+        ),
+    ):
+        await hass.config_entries.async_setup(entry.entry_id)
+        await hass.async_block_till_done()
+
+    assert entry.state is ConfigEntryState.SETUP_RETRY
+
+
+async def test_oauth_client_error_is_retryable(hass: HomeAssistant) -> None:
+    """Test OAuth transport client errors mark entry for setup retry."""
+    entry = await mock_august_config_entry(hass)
+
+    with patch(
+        "homeassistant.helpers.config_entry_oauth2_flow.OAuth2Session.async_ensure_token_valid",
+        side_effect=ClientError("connection error"),
+    ):
+        await hass.config_entries.async_setup(entry.entry_id)
+        await hass.async_block_till_done()
+
+    assert entry.state is ConfigEntryState.SETUP_RETRY