Migrate google_mail OAuth token refresh exception handling (#165371)

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

homeassistant/components/google_mail/api.py

@@ -2,8 +2,7 @@
 
 from functools import partial
 
-from aiohttp.client_exceptions import ClientError, ClientResponseError
-from google.auth.exceptions import RefreshError
+from aiohttp.client_exceptions import ClientError
 from google.oauth2.credentials import Credentials
 from googleapiclient.discovery import Resource, build
 
@@ -14,6 +13,8 @@ from homeassistant.exceptions import (
     ConfigEntryAuthFailed,
     ConfigEntryNotReady,
     HomeAssistantError,
+    OAuth2TokenRequestError,
+    OAuth2TokenRequestReauthError,
 )
 from homeassistant.helpers import config_entry_oauth2_flow
 
@@ -37,24 +38,26 @@ class AsyncConfigEntryAuth:
 
     async def check_and_refresh_token(self) -> str:
         """Check the token."""
+        setup_in_progress = (
+            self.oauth_session.config_entry.state is ConfigEntryState.SETUP_IN_PROGRESS
+        )
+
         try:
             await self.oauth_session.async_ensure_token_valid()
-        except (RefreshError, ClientResponseError, ClientError) as ex:
-            if (
-                self.oauth_session.config_entry.state
-                is ConfigEntryState.SETUP_IN_PROGRESS
-            ):
-                if isinstance(ex, ClientResponseError) and 400 <= ex.status < 500:
-                    raise ConfigEntryAuthFailed(
-                        "OAuth session is not valid, reauth required"
-                    ) from ex
+        except OAuth2TokenRequestReauthError as ex:
+            if setup_in_progress:
+                raise ConfigEntryAuthFailed(
+                    "OAuth session is not valid, reauth required"
+                ) from ex
+            self.oauth_session.config_entry.async_start_reauth(self.oauth_session.hass)
+            raise
+        except OAuth2TokenRequestError as ex:
+            if setup_in_progress:
+                raise ConfigEntryNotReady from ex
+            raise
+        except ClientError as ex:
+            if setup_in_progress:
                 raise ConfigEntryNotReady from ex
-            if isinstance(ex, RefreshError) or (
-                hasattr(ex, "status") and ex.status == 400
-            ):
-                self.oauth_session.config_entry.async_start_reauth(
-                    self.oauth_session.hass
-                )
             raise HomeAssistantError(ex) from ex
         return self.access_token
 

tests/components/google_mail/test_init.py

@@ -2,14 +2,19 @@
 
 import http
 import time
-from unittest.mock import patch
+from unittest.mock import Mock, patch
 
 from aiohttp.client_exceptions import ClientError
 import pytest
 
 from homeassistant.components.google_mail import DOMAIN
-from homeassistant.config_entries import ConfigEntryState
+from homeassistant.config_entries import SOURCE_REAUTH, ConfigEntryState
 from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import (
+    OAuth2TokenRequestError,
+    OAuth2TokenRequestReauthError,
+    OAuth2TokenRequestTransientError,
+)
 from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.config_entry_oauth2_flow import (
     ImplementationUnavailableError,
@@ -76,13 +81,23 @@ async def test_expired_token_refresh_success(
             http.HTTPStatus.INTERNAL_SERVER_ERROR,
             ConfigEntryState.SETUP_RETRY,
         ),
+        (
+            time.time() - 3600,
+            http.HTTPStatus.TOO_MANY_REQUESTS,
+            ConfigEntryState.SETUP_RETRY,
+        ),
         (
             time.time() - 3600,
             http.HTTPStatus.BAD_REQUEST,
             ConfigEntryState.SETUP_ERROR,
         ),
     ],
-    ids=["failure_requires_reauth", "transient_failure", "revoked_auth"],
+    ids=[
+        "failure_requires_reauth",
+        "transient_failure",
+        "rate_limited",
+        "revoked_auth",
+    ],
 )
 async def test_expired_token_refresh_failure(
     hass: HomeAssistant,
@@ -123,6 +138,69 @@ async def test_expired_token_refresh_client_error(
     assert entries[0].state is ConfigEntryState.SETUP_RETRY
 
 
+async def test_token_refresh_reauth_error_during_setup(
+    hass: HomeAssistant,
+    setup_integration: ComponentSetup,
+) -> None:
+    """Test setup starts reauth for OAuth reauth errors."""
+    with patch(
+        "homeassistant.components.google_mail.OAuth2Session.async_ensure_token_valid",
+        side_effect=OAuth2TokenRequestReauthError(
+            request_info=Mock(),
+            domain=DOMAIN,
+        ),
+    ):
+        await setup_integration()
+
+    entry = hass.config_entries.async_entries(DOMAIN)[0]
+    assert entry.state is ConfigEntryState.SETUP_ERROR
+
+    flows = hass.config_entries.flow.async_progress()
+    assert len(flows) == 1
+    flow = flows[0]
+    assert flow["step_id"] == "reauth_confirm"
+    assert flow["handler"] == DOMAIN
+    assert flow["context"]["source"] == SOURCE_REAUTH
+
+
+async def test_token_refresh_transient_error_during_setup(
+    hass: HomeAssistant,
+    setup_integration: ComponentSetup,
+) -> None:
+    """Test setup retries for transient OAuth token refresh errors."""
+    with patch(
+        "homeassistant.components.google_mail.OAuth2Session.async_ensure_token_valid",
+        side_effect=OAuth2TokenRequestTransientError(
+            request_info=Mock(),
+            domain=DOMAIN,
+        ),
+    ):
+        await setup_integration()
+
+    entry = hass.config_entries.async_entries(DOMAIN)[0]
+    assert entry.state is ConfigEntryState.SETUP_RETRY
+    assert not hass.config_entries.flow.async_progress()
+
+
+async def test_token_refresh_error_during_setup(
+    hass: HomeAssistant,
+    setup_integration: ComponentSetup,
+) -> None:
+    """Test generic OAuth token refresh errors retry setup."""
+    with patch(
+        "homeassistant.components.google_mail.OAuth2Session.async_ensure_token_valid",
+        side_effect=OAuth2TokenRequestError(
+            request_info=Mock(),
+            domain=DOMAIN,
+        ),
+    ):
+        await setup_integration()
+
+    entry = hass.config_entries.async_entries(DOMAIN)[0]
+    assert entry.state is ConfigEntryState.SETUP_RETRY
+    assert not hass.config_entries.flow.async_progress()
+
+
 async def test_device_info(
     hass: HomeAssistant,
     device_registry: dr.DeviceRegistry,

tests/components/google_mail/test_sensor.py

@@ -1,9 +1,9 @@
 """Sensor tests for the Google Mail integration."""
 
 from datetime import timedelta
-from unittest.mock import patch
+from unittest.mock import Mock, patch
 
-from google.auth.exceptions import RefreshError
+from aiohttp.client_exceptions import ClientResponseError
 from httplib2 import Response
 import pytest
 
@@ -12,6 +12,11 @@ from homeassistant.components.google_mail.const import DOMAIN
 from homeassistant.components.sensor import SensorDeviceClass
 from homeassistant.const import ATTR_DEVICE_CLASS, STATE_UNKNOWN
 from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import (
+    OAuth2TokenRequestError,
+    OAuth2TokenRequestReauthError,
+    OAuth2TokenRequestTransientError,
+)
 from homeassistant.util import dt as dt_util
 
 from .conftest import SENSOR, TOKEN, ComponentSetup
@@ -56,12 +61,19 @@ async def test_sensors(
 
 
 async def test_sensor_reauth_trigger(
-    hass: HomeAssistant, setup_integration: ComponentSetup
+    hass: HomeAssistant,
+    setup_integration: ComponentSetup,
 ) -> None:
     """Test reauth is triggered after a refresh error."""
     await setup_integration()
 
-    with patch(TOKEN, side_effect=RefreshError):
+    with patch(
+        TOKEN,
+        side_effect=OAuth2TokenRequestReauthError(
+            request_info=Mock(),
+            domain=DOMAIN,
+        ),
+    ):
         next_update = dt_util.utcnow() + timedelta(minutes=15)
         async_fire_time_changed(hass, next_update)
         await hass.async_block_till_done(wait_background_tasks=True)
@@ -73,3 +85,34 @@ async def test_sensor_reauth_trigger(
     assert flow["step_id"] == "reauth_confirm"
     assert flow["handler"] == DOMAIN
     assert flow["context"]["source"] == config_entries.SOURCE_REAUTH
+
+
+@pytest.mark.parametrize(
+    "side_effect",
+    [
+        OAuth2TokenRequestTransientError(request_info=Mock(), domain=DOMAIN),
+        OAuth2TokenRequestError(request_info=Mock(), domain=DOMAIN),
+        ClientResponseError(request_info=Mock(), history=(), status=401),
+        ClientResponseError(request_info=Mock(), history=(), status=429),
+    ],
+    ids=[
+        "oauth_transient_error",
+        "oauth_generic_error",
+        "legacy_client_response_4xx",
+        "legacy_rate_limited",
+    ],
+)
+async def test_sensor_token_error_no_reauth(
+    hass: HomeAssistant,
+    setup_integration: ComponentSetup,
+    side_effect: Exception | type[Exception],
+) -> None:
+    """Test retryable/runtime token errors do not start reauth."""
+    await setup_integration()
+
+    with patch(TOKEN, side_effect=side_effect):
+        next_update = dt_util.utcnow() + timedelta(minutes=15)
+        async_fire_time_changed(hass, next_update)
+        await hass.async_block_till_done(wait_background_tasks=True)
+
+    assert not hass.config_entries.flow.async_progress()

tests/components/google_mail/test_services.py

@@ -1,15 +1,13 @@
 """Services tests for the Google Mail integration."""
 
-from unittest.mock import patch
+from unittest.mock import Mock, patch
 
-from aiohttp.client_exceptions import ClientResponseError
-from google.auth.exceptions import RefreshError
 import pytest
 
 from homeassistant import config_entries
 from homeassistant.components.google_mail import DOMAIN
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import HomeAssistantError
+from homeassistant.exceptions import HomeAssistantError, OAuth2TokenRequestReauthError
 
 from .conftest import BUILD, SENSOR, TOKEN, ComponentSetup
 
@@ -60,22 +58,23 @@ async def test_set_vacation(
     assert len(mock_client.mock_calls) == 5
 
 
-@pytest.mark.parametrize(
-    ("side_effect"),
-    [
-        (RefreshError,),
-        (ClientResponseError("", (), status=400),),
-    ],
-)
 async def test_reauth_trigger(
     hass: HomeAssistant,
     setup_integration: ComponentSetup,
-    side_effect,
 ) -> None:
     """Test reauth is triggered after a refresh error during service call."""
     await setup_integration()
 
-    with patch(TOKEN, side_effect=side_effect), pytest.raises(HomeAssistantError):
+    with (
+        patch(
+            TOKEN,
+            side_effect=OAuth2TokenRequestReauthError(
+                request_info=Mock(),
+                domain=DOMAIN,
+            ),
+        ),
+        pytest.raises(HomeAssistantError),
+    ):
         await hass.services.async_call(
             DOMAIN,
             "set_vacation",