Add reauth support in fitbit (home-assistant#101178)

* Add reauth support in fitbit * Update tests/components/fitbit/test_config_flow.py Co-authored-by: Martin Hjelmare <[email protected]> * Improve http status error code handling --------- Co-authored-by: Martin Hjelmare <[email protected]>
Caius-Bonus · Oct 2, 2023 · 4e4b8de · 4e4b8de
1 parent 4c24ff6
commit 4e4b8de
Show file tree

Hide file tree

Showing 5 changed files with 220 additions and 5 deletions.
diff --git a/homeassistant/components/fitbit/__init__.py b/homeassistant/components/fitbit/__init__.py
@@ -1,11 +1,13 @@
 """The fitbit component."""
 
+from http import HTTPStatus
+
 import aiohttp
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import ConfigEntryNotReady
+from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
 from homeassistant.helpers import config_entry_oauth2_flow
 
 from . import api
@@ -29,6 +31,10 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     )
     try:
         await fitbit_api.async_get_access_token()
+    except aiohttp.ClientResponseError as err:
+        if err.status == HTTPStatus.UNAUTHORIZED:
+            raise ConfigEntryAuthFailed from err
+        raise ConfigEntryNotReady from err
     except aiohttp.ClientError as err:
         raise ConfigEntryNotReady from err
 

diff --git a/homeassistant/components/fitbit/config_flow.py b/homeassistant/components/fitbit/config_flow.py
@@ -1,10 +1,12 @@
 """Config flow for fitbit."""
 
+from collections.abc import Mapping
 import logging
 from typing import Any
 
 from fitbit.exceptions import HTTPException
 
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_TOKEN
 from homeassistant.data_entry_flow import FlowResult
 from homeassistant.helpers import config_entry_oauth2_flow
@@ -22,6 +24,8 @@ class OAuth2FlowHandler(
 
     DOMAIN = DOMAIN
 
+    reauth_entry: ConfigEntry | None = None
+
     @property
     def logger(self) -> logging.Logger:
         """Return logger."""
@@ -32,9 +36,24 @@ def extra_authorize_data(self) -> dict[str, str]:
         """Extra data that needs to be appended to the authorize url."""
         return {
             "scope": " ".join(OAUTH_SCOPES),
-            "prompt": "consent",
+            "prompt": "consent" if not self.reauth_entry else "none",
         }
 
+    async def async_step_reauth(self, entry_data: Mapping[str, Any]) -> FlowResult:
+        """Perform reauth upon an API authentication error."""
+        self.reauth_entry = self.hass.config_entries.async_get_entry(
+            self.context["entry_id"]
+        )
+        return await self.async_step_reauth_confirm()
+
+    async def async_step_reauth_confirm(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
+        """Confirm reauth dialog."""
+        if user_input is None:
+            return self.async_show_form(step_id="reauth_confirm")
+        return await self.async_step_user()
+
     async def async_oauth_create_entry(self, data: dict[str, Any]) -> FlowResult:
         """Create an entry for the flow, or update existing entry."""
 
@@ -45,6 +64,13 @@ async def async_oauth_create_entry(self, data: dict[str, Any]) -> FlowResult:
             _LOGGER.error("Failed to fetch user profile for Fitbit API: %s", err)
             return self.async_abort(reason="cannot_connect")
 
+        if self.reauth_entry:
+            if self.reauth_entry.unique_id != profile.encoded_id:
+                return self.async_abort(reason="wrong_account")
+            self.hass.config_entries.async_update_entry(self.reauth_entry, data=data)
+            await self.hass.config_entries.async_reload(self.reauth_entry.entry_id)
+            return self.async_abort(reason="reauth_successful")
+
         await self.async_set_unique_id(profile.encoded_id)
         self._abort_if_unique_id_configured()
         return self.async_create_entry(title=profile.full_name, data=data)

diff --git a/homeassistant/components/fitbit/strings.json b/homeassistant/components/fitbit/strings.json
@@ -6,6 +6,10 @@
       },
       "auth": {
         "title": "Link Fitbit"
+      },
+      "reauth_confirm": {
+        "title": "[%key:common::config_flow::title::reauth%]",
+        "description": "The Fitbit integration needs to re-authenticate your account"
       }
     },
     "abort": {
@@ -15,7 +19,9 @@
       "oauth_error": "[%key:common::config_flow::abort::oauth2_error%]",
       "missing_configuration": "[%key:common::config_flow::abort::oauth2_missing_configuration%]",
       "invalid_access_token": "[%key:common::config_flow::error::invalid_access_token%]",
-      "unknown": "[%key:common::config_flow::error::unknown%]"
+      "reauth_successful": "[%key:common::config_flow::abort::reauth_successful%]",
+      "unknown": "[%key:common::config_flow::error::unknown%]",
+      "wrong_account": "The user credentials provided do not match this Fitbit account."
     },
     "create_entry": {
       "default": "[%key:common::config_flow::create_entry::authenticated%]"

diff --git a/tests/components/fitbit/test_config_flow.py b/tests/components/fitbit/test_config_flow.py
@@ -4,6 +4,7 @@
 from http import HTTPStatus
 from unittest.mock import patch
 
+import pytest
 from requests_mock.mocker import Mocker
 
 from homeassistant import config_entries
@@ -313,3 +314,148 @@ async def test_platform_setup_without_import(
     issue = issue_registry.issues.get((DOMAIN, "deprecated_yaml"))
     assert issue
     assert issue.translation_key == "deprecated_yaml_no_import"
+
+
+async def test_reauth_flow(
+    hass: HomeAssistant,
+    config_entry: MockConfigEntry,
+    hass_client_no_auth: ClientSessionGenerator,
+    aioclient_mock: AiohttpClientMocker,
+    current_request_with_host: None,
+    profile: None,
+    setup_credentials: None,
+) -> None:
+    """Test OAuth reauthentication flow will update existing config entry."""
+    config_entry.add_to_hass(hass)
+
+    entries = hass.config_entries.async_entries(DOMAIN)
+    assert len(entries) == 1
+
+    # config_entry.req initiates reauth
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        context={
+            "source": config_entries.SOURCE_REAUTH,
+            "entry_id": config_entry.entry_id,
+        },
+    )
+    assert result["type"] == "form"
+    assert result["step_id"] == "reauth_confirm"
+
+    result = await hass.config_entries.flow.async_configure(
+        flow_id=result["flow_id"],
+        user_input={},
+    )
+    assert result["type"] == FlowResultType.EXTERNAL_STEP
+    state = config_entry_oauth2_flow._encode_jwt(
+        hass,
+        {
+            "flow_id": result["flow_id"],
+            "redirect_uri": REDIRECT_URL,
+        },
+    )
+    assert result["url"] == (
+        f"{OAUTH2_AUTHORIZE}?response_type=code&client_id={CLIENT_ID}"
+        f"&redirect_uri={REDIRECT_URL}"
+        f"&state={state}"
+        "&scope=activity+heartrate+nutrition+profile+settings+sleep+weight&prompt=none"
+    )
+
+    client = await hass_client_no_auth()
+    resp = await client.get(f"/auth/external/callback?code=abcd&state={state}")
+    assert resp.status == 200
+    assert resp.headers["content-type"] == "text/html; charset=utf-8"
+
+    aioclient_mock.post(
+        OAUTH2_TOKEN,
+        json={
+            "refresh_token": "updated-refresh-token",
+            "access_token": "updated-access-token",
+            "type": "Bearer",
+            "expires_in": 60,
+        },
+    )
+
+    with patch(
+        "homeassistant.components.fitbit.async_setup_entry", return_value=True
+    ) as mock_setup:
+        result = await hass.config_entries.flow.async_configure(result["flow_id"])
+
+    assert result.get("type") == FlowResultType.ABORT
+    assert result.get("reason") == "reauth_successful"
+
+    assert len(hass.config_entries.async_entries(DOMAIN)) == 1
+    assert len(mock_setup.mock_calls) == 1
+
+    assert config_entry.data["token"]["refresh_token"] == "updated-refresh-token"
+
+
+@pytest.mark.parametrize("profile_id", ["other-user-id"])
+async def test_reauth_wrong_user_id(
+    hass: HomeAssistant,
+    config_entry: MockConfigEntry,
+    hass_client_no_auth: ClientSessionGenerator,
+    aioclient_mock: AiohttpClientMocker,
+    current_request_with_host: None,
+    profile: None,
+    setup_credentials: None,
+) -> None:
+    """Test OAuth reauthentication where the wrong user is selected."""
+    config_entry.add_to_hass(hass)
+
+    entries = hass.config_entries.async_entries(DOMAIN)
+    assert len(entries) == 1
+
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        context={
+            "source": config_entries.SOURCE_REAUTH,
+            "entry_id": config_entry.entry_id,
+        },
+    )
+    assert result["type"] == FlowResultType.FORM
+    assert result["step_id"] == "reauth_confirm"
+
+    result = await hass.config_entries.flow.async_configure(
+        flow_id=result["flow_id"],
+        user_input={},
+    )
+    assert result["type"] == FlowResultType.EXTERNAL_STEP
+    state = config_entry_oauth2_flow._encode_jwt(
+        hass,
+        {
+            "flow_id": result["flow_id"],
+            "redirect_uri": REDIRECT_URL,
+        },
+    )
+    assert result["url"] == (
+        f"{OAUTH2_AUTHORIZE}?response_type=code&client_id={CLIENT_ID}"
+        f"&redirect_uri={REDIRECT_URL}"
+        f"&state={state}"
+        "&scope=activity+heartrate+nutrition+profile+settings+sleep+weight&prompt=none"
+    )
+
+    client = await hass_client_no_auth()
+    resp = await client.get(f"/auth/external/callback?code=abcd&state={state}")
+    assert resp.status == 200
+    assert resp.headers["content-type"] == "text/html; charset=utf-8"
+
+    aioclient_mock.post(
+        OAUTH2_TOKEN,
+        json={
+            "refresh_token": "updated-refresh-token",
+            "access_token": "updated-access-token",
+            "type": "Bearer",
+            "expires_in": 60,
+        },
+    )
+
+    with patch(
+        "homeassistant.components.fitbit.async_setup_entry", return_value=True
+    ) as mock_setup:
+        result = await hass.config_entries.flow.async_configure(result["flow_id"])
+
+    assert result.get("type") == FlowResultType.ABORT
+    assert result.get("reason") == "wrong_account"
+
+    assert len(mock_setup.mock_calls) == 0
diff --git a/tests/components/fitbit/test_init.py b/tests/components/fitbit/test_init.py
@@ -43,18 +43,26 @@ async def test_setup(
     assert config_entry.state == ConfigEntryState.NOT_LOADED
 
 
-@pytest.mark.parametrize("token_expiration_time", [12345])
+@pytest.mark.parametrize(
+    ("token_expiration_time", "server_status"),
+    [
+        (12345, HTTPStatus.INTERNAL_SERVER_ERROR),
+        (12345, HTTPStatus.FORBIDDEN),
+        (12345, HTTPStatus.NOT_FOUND),
+    ],
+)
 async def test_token_refresh_failure(
     integration_setup: Callable[[], Awaitable[bool]],
     config_entry: MockConfigEntry,
     aioclient_mock: AiohttpClientMocker,
     setup_credentials: None,
+    server_status: HTTPStatus,
 ) -> None:
     """Test where token is expired and the refresh attempt fails and will be retried."""
 
     aioclient_mock.post(
         OAUTH2_TOKEN,
-        status=HTTPStatus.INTERNAL_SERVER_ERROR,
+        status=server_status,
     )
 
     assert not await integration_setup()
@@ -94,3 +102,26 @@ async def test_token_refresh_success(
         config_entry.data["token"]["access_token"]
         == SERVER_ACCESS_TOKEN["access_token"]
     )
+
+
+@pytest.mark.parametrize("token_expiration_time", [12345])
+async def test_token_requires_reauth(
+    hass: HomeAssistant,
+    integration_setup: Callable[[], Awaitable[bool]],
+    config_entry: MockConfigEntry,
+    aioclient_mock: AiohttpClientMocker,
+    setup_credentials: None,
+) -> None:
+    """Test where token is expired and the refresh attempt requires reauth."""
+
+    aioclient_mock.post(
+        OAUTH2_TOKEN,
+        status=HTTPStatus.UNAUTHORIZED,
+    )
+
+    assert not await integration_setup()
+    assert config_entry.state == ConfigEntryState.SETUP_ERROR
+
+    flows = hass.config_entries.flow.async_progress()
+    assert len(flows) == 1
+    assert flows[0]["step_id"] == "reauth_confirm"