Fix onedrive oauth errors during setup (#179398)

homeassistant/components/onedrive/__init__.py

@@ -15,7 +15,12 @@ from onedrive_personal_sdk.exceptions import (
 
 from homeassistant.const import CONF_ACCESS_TOKEN, Platform
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
+from homeassistant.exceptions import (
+    ConfigEntryAuthFailed,
+    ConfigEntryNotReady,
+    OAuth2TokenRequestError,
+    OAuth2TokenRequestReauthError,
+)
 from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.config_entry_oauth2_flow import (
@@ -175,6 +180,18 @@ async def _get_onedrive_client(
         ) from err
     session = OAuth2Session(hass, entry, implementation)
 
+    # Refresh up front, so a failure surfaces here instead of from inside the client
+    try:
+        await session.async_ensure_token_valid()
+    except OAuth2TokenRequestReauthError as err:
+        raise ConfigEntryAuthFailed(
+            translation_domain=DOMAIN, translation_key="authentication_failed"
+        ) from err
+    except OAuth2TokenRequestError as err:
+        raise ConfigEntryNotReady(
+            translation_domain=DOMAIN, translation_key="connection_error"
+        ) from err
+
     async def get_access_token() -> str:
         await session.async_ensure_token_valid()
         return cast(str, session.token[CONF_ACCESS_TOKEN])

homeassistant/components/onedrive_for_business/__init__.py

@@ -13,7 +13,12 @@ from onedrive_personal_sdk.exceptions import (
 
 from homeassistant.const import CONF_ACCESS_TOKEN, Platform
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
+from homeassistant.exceptions import (
+    ConfigEntryAuthFailed,
+    ConfigEntryNotReady,
+    OAuth2TokenRequestError,
+    OAuth2TokenRequestReauthError,
+)
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.config_entry_oauth2_flow import (
     ImplementationUnavailableError,
@@ -100,6 +105,18 @@ async def _get_onedrive_client(
             ) from err
     session = OAuth2Session(hass, entry, implementation)
 
+    # Refresh up front, so a failure surfaces here instead of from inside the client
+    try:
+        await session.async_ensure_token_valid()
+    except OAuth2TokenRequestReauthError as err:
+        raise ConfigEntryAuthFailed(
+            translation_domain=DOMAIN, translation_key="authentication_failed"
+        ) from err
+    except OAuth2TokenRequestError as err:
+        raise ConfigEntryNotReady(
+            translation_domain=DOMAIN, translation_key="connection_error"
+        ) from err
+
     async def get_access_token() -> str:
         await session.async_ensure_token_valid()
         return cast(str, session.token[CONF_ACCESS_TOKEN])

homeassistant/components/onedrive_for_business/strings.json

@@ -96,6 +96,9 @@
     "authentication_failed": {
       "message": "[%key:component::onedrive::exceptions::authentication_failed::message%]"
     },
+    "connection_error": {
+      "message": "[%key:component::onedrive::config::abort::connection_error%]"
+    },
     "failed_to_get_folder": {
       "message": "[%key:component::onedrive::exceptions::failed_to_get_folder::message%]"
     },

tests/components/onedrive/test_init.py

@@ -2,6 +2,7 @@
 
 from copy import copy
 from html import escape
+from http import HTTPStatus
 from json import dumps
 from unittest.mock import MagicMock, patch
 
@@ -19,6 +20,7 @@ from homeassistant.components.onedrive.const import (
     CONF_FOLDER_ID,
     CONF_FOLDER_NAME,
     DOMAIN,
+    OAUTH2_TOKEN,
 )
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import HomeAssistant
@@ -31,6 +33,7 @@ from . import setup_integration
 from .const import BACKUP_METADATA
 
 from tests.common import MockConfigEntry
+from tests.test_util.aiohttp import AiohttpClientMocker
 
 
 async def test_load_unload_config_entry(
@@ -58,6 +61,54 @@ async def test_load_unload_config_entry(
     assert mock_config_entry.state is ConfigEntryState.NOT_LOADED
 
 
+@pytest.mark.parametrize(
+    ("status", "state", "reason", "reauth_expected"),
+    [
+        pytest.param(
+            HTTPStatus.BAD_REQUEST,
+            ConfigEntryState.SETUP_ERROR,
+            "Authentication failed",
+            True,
+            id="reauth",
+        ),
+        pytest.param(
+            HTTPStatus.TOO_MANY_REQUESTS,
+            ConfigEntryState.SETUP_RETRY,
+            "Failed to connect to OneDrive",
+            False,
+            id="transient",
+        ),
+        pytest.param(
+            HTTPStatus.INTERNAL_SERVER_ERROR,
+            ConfigEntryState.SETUP_RETRY,
+            "Failed to connect to OneDrive",
+            False,
+            id="server_error",
+        ),
+    ],
+)
+@pytest.mark.parametrize("expires_at", [0], ids=["expired"])
+async def test_token_refresh_errors(
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    aioclient_mock: AiohttpClientMocker,
+    status: HTTPStatus,
+    state: ConfigEntryState,
+    reason: str,
+    reauth_expected: bool,
+) -> None:
+    """Test a failing token refresh during setup."""
+    aioclient_mock.post(OAUTH2_TOKEN, status=status, json={})
+    mock_config_entry.add_to_hass(hass)
+
+    assert not await hass.config_entries.async_setup(mock_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert mock_config_entry.state is state
+    assert mock_config_entry.reason == reason
+    assert bool(hass.config_entries.flow.async_progress()) is reauth_expected
+
+
 @pytest.mark.parametrize(
     ("side_effect", "state"),
     [

tests/components/onedrive_for_business/test_init.py

@@ -1,6 +1,7 @@
 """Test the OneDrive setup."""
 
 from copy import copy
+from http import HTTPStatus
 from unittest.mock import MagicMock, patch
 
 from onedrive_personal_sdk.exceptions import (
@@ -14,6 +15,8 @@ import pytest
 from homeassistant.components.onedrive_for_business.const import (
     CONF_FOLDER_ID,
     CONF_FOLDER_PATH,
+    CONF_TENANT_ID,
+    OAUTH2_TOKEN,
 )
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import HomeAssistant
@@ -24,6 +27,7 @@ from homeassistant.helpers.config_entry_oauth2_flow import (
 from . import setup_integration
 
 from tests.common import MockConfigEntry
+from tests.test_util.aiohttp import AiohttpClientMocker
 
 
 async def test_load_unload_config_entry(
@@ -51,6 +55,58 @@ async def test_load_unload_config_entry(
     assert mock_config_entry.state is ConfigEntryState.NOT_LOADED
 
 
+@pytest.mark.parametrize(
+    ("status", "state", "reason", "reauth_expected"),
+    [
+        pytest.param(
+            HTTPStatus.BAD_REQUEST,
+            ConfigEntryState.SETUP_ERROR,
+            "Authentication failed",
+            True,
+            id="reauth",
+        ),
+        pytest.param(
+            HTTPStatus.TOO_MANY_REQUESTS,
+            ConfigEntryState.SETUP_RETRY,
+            "Failed to connect to OneDrive",
+            False,
+            id="transient",
+        ),
+        pytest.param(
+            HTTPStatus.INTERNAL_SERVER_ERROR,
+            ConfigEntryState.SETUP_RETRY,
+            "Failed to connect to OneDrive",
+            False,
+            id="server_error",
+        ),
+    ],
+)
+@pytest.mark.parametrize("expires_at", [0], ids=["expired"])
+async def test_token_refresh_errors(
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    aioclient_mock: AiohttpClientMocker,
+    status: HTTPStatus,
+    state: ConfigEntryState,
+    reason: str,
+    reauth_expected: bool,
+) -> None:
+    """Test a failing token refresh during setup."""
+    aioclient_mock.post(
+        OAUTH2_TOKEN.format(tenant_id=mock_config_entry.data[CONF_TENANT_ID]),
+        status=status,
+        json={},
+    )
+    mock_config_entry.add_to_hass(hass)
+
+    assert not await hass.config_entries.async_setup(mock_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert mock_config_entry.state is state
+    assert mock_config_entry.reason == reason
+    assert bool(hass.config_entries.flow.async_progress()) is reauth_expected
+
+
 @pytest.mark.parametrize(
     ("side_effect", "state"),
     [