Add reauth support to LaMetric (home-assistant#80355)

* Add reauth support to LaMetric

* Adjust docblock

homeassistant/components/lametric/config_flow.py

@@ -1,6 +1,7 @@
 """Config flow to configure the LaMetric integration."""
 from __future__ import annotations
 
+from collections.abc import Mapping
 from ipaddress import ip_address
 import logging
 from typing import Any
@@ -27,6 +28,7 @@
     ATTR_UPNP_SERIAL,
     SsdpServiceInfo,
 )
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_API_KEY, CONF_DEVICE, CONF_HOST, CONF_MAC
 from homeassistant.data_entry_flow import AbortFlow, FlowResult
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
@@ -56,6 +58,7 @@ class LaMetricFlowHandler(AbstractOAuth2FlowHandler, domain=DOMAIN):
     discovered_host: str
     discovered_serial: str
     discovered: bool = False
+    reauth_entry: ConfigEntry | None = None
 
     @property
     def logger(self) -> logging.Logger:
@@ -103,6 +106,13 @@ async def async_step_ssdp(self, discovery_info: SsdpServiceInfo) -> FlowResult:
         self.discovered_serial = serial
         return await self.async_step_choice_enter_manual_or_fetch_cloud()
 
+    async def async_step_reauth(self, entry_data: Mapping[str, Any]) -> FlowResult:
+        """Handle initiation of re-authentication with LaMetric."""
+        self.reauth_entry = self.hass.config_entries.async_get_entry(
+            self.context["entry_id"]
+        )
+        return await self.async_step_choice_enter_manual_or_fetch_cloud()
+
     async def async_step_choice_enter_manual_or_fetch_cloud(
         self, user_input: dict[str, Any] | None = None
     ) -> FlowResult:
@@ -120,6 +130,8 @@ async def async_step_manual_entry(
         if user_input is not None:
             if self.discovered:
                 host = self.discovered_host
+            elif self.reauth_entry:
+                host = self.reauth_entry.data[CONF_HOST]
             else:
                 host = user_input[CONF_HOST]
 
@@ -142,7 +154,7 @@ async def async_step_manual_entry(
                 TextSelectorConfig(type=TextSelectorType.PASSWORD)
             )
         }
-        if not self.discovered:
+        if not self.discovered and not self.reauth_entry:
             schema = {vol.Required(CONF_HOST): TextSelector()} | schema
 
         return self.async_show_form(
@@ -173,6 +185,10 @@ async def async_step_cloud_select_device(
         """Handle device selection from devices offered by the cloud."""
         if self.discovered:
             user_input = {CONF_DEVICE: self.discovered_serial}
+        elif self.reauth_entry:
+            if self.reauth_entry.unique_id not in self.devices:
+                return self.async_abort(reason="reauth_device_not_found")
+            user_input = {CONF_DEVICE: self.reauth_entry.unique_id}
         elif len(self.devices) == 1:
             user_input = {CONF_DEVICE: list(self.devices.values())[0].serial_number}
 
@@ -223,10 +239,11 @@ async def _async_step_create_entry(self, host: str, api_key: str) -> FlowResult:
 
         device = await lametric.device()
 
-        await self.async_set_unique_id(device.serial_number)
-        self._abort_if_unique_id_configured(
-            updates={CONF_HOST: lametric.host, CONF_API_KEY: lametric.api_key}
-        )
+        if not self.reauth_entry:
+            await self.async_set_unique_id(device.serial_number)
+            self._abort_if_unique_id_configured(
+                updates={CONF_HOST: lametric.host, CONF_API_KEY: lametric.api_key}
+            )
 
         await lametric.notify(
             notification=Notification(
@@ -240,6 +257,20 @@ async def _async_step_create_entry(self, host: str, api_key: str) -> FlowResult:
             )
         )
 
+        if self.reauth_entry:
+            self.hass.config_entries.async_update_entry(
+                self.reauth_entry,
+                data={
+                    **self.reauth_entry.data,
+                    CONF_HOST: lametric.host,
+                    CONF_API_KEY: lametric.api_key,
+                },
+            )
+            self.hass.async_create_task(
+                self.hass.config_entries.async_reload(self.reauth_entry.entry_id)
+            )
+            return self.async_abort(reason="reauth_successful")
+
         return self.async_create_entry(
             title=device.name,
             data={

homeassistant/components/lametric/coordinator.py

@@ -1,11 +1,12 @@
 """DataUpdateCoordinator for the LaMatric integration."""
 from __future__ import annotations
 
-from demetriek import Device, LaMetricDevice, LaMetricError
+from demetriek import Device, LaMetricAuthenticationError, LaMetricDevice, LaMetricError
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_API_KEY, CONF_HOST
 from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import ConfigEntryAuthFailed
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
@@ -32,6 +33,8 @@ async def _async_update_data(self) -> Device:
         """Fetch device information of the LaMetric device."""
         try:
             return await self.lametric.device()
+        except LaMetricAuthenticationError as err:
+            raise ConfigEntryAuthFailed from err
         except LaMetricError as ex:
             raise UpdateFailed(
                 "Could not fetch device information from LaMetric device"

homeassistant/components/lametric/strings.json

@@ -39,6 +39,8 @@
       "missing_configuration": "The LaMetric integration is not configured. Please follow the documentation.",
       "no_devices": "The authorized user has no LaMetric devices",
       "no_url_available": "[%key:common::config_flow::abort::oauth2_no_url_available%]",
+      "reauth_device_not_found": "The device you are trying to re-authenticate is not found in this LaMetric account",
+      "reauth_successful": "[%key:common::config_flow::abort::reauth_successful%]",
       "unknown": "[%key:common::config_flow::error::unknown%]"
     }
   },

homeassistant/components/lametric/translations/en.json

@@ -8,6 +8,8 @@
             "missing_configuration": "The LaMetric integration is not configured. Please follow the documentation.",
             "no_devices": "The authorized user has no LaMetric devices",
             "no_url_available": "No URL available. For information about this error, [check the help section]({docs_url})",
+            "reauth_device_not_found": "The device you are trying to re-authenticate is not found in this LaMetric account",
+            "reauth_successful": "Re-authentication was successful",
             "unknown": "Unexpected error"
         },
         "error": {

tests/components/lametric/test_config_flow.py

@@ -18,7 +18,12 @@
     ATTR_UPNP_SERIAL,
     SsdpServiceInfo,
 )
-from homeassistant.config_entries import SOURCE_DHCP, SOURCE_SSDP, SOURCE_USER
+from homeassistant.config_entries import (
+    SOURCE_DHCP,
+    SOURCE_REAUTH,
+    SOURCE_SSDP,
+    SOURCE_USER,
+)
 from homeassistant.const import CONF_API_KEY, CONF_DEVICE, CONF_HOST, CONF_MAC
 from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResultType
@@ -743,3 +748,173 @@ async def test_dhcp_unknown_device(
 
     assert result.get("type") == FlowResultType.ABORT
     assert result.get("reason") == "unknown"
+
+
+async def test_reauth_cloud_import(
+    hass: HomeAssistant,
+    hass_client_no_auth: Callable[[], Awaitable[TestClient]],
+    aioclient_mock: AiohttpClientMocker,
+    current_request_with_host: None,
+    mock_setup_entry: MagicMock,
+    mock_lametric_cloud_config_flow: MagicMock,
+    mock_lametric_config_flow: MagicMock,
+    mock_config_entry: MockConfigEntry,
+) -> None:
+    """Test reauth flow importing api keys from the cloud."""
+    mock_config_entry.add_to_hass(hass)
+
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        context={
+            "source": SOURCE_REAUTH,
+            "unique_id": mock_config_entry.unique_id,
+            "entry_id": mock_config_entry.entry_id,
+        },
+        data=mock_config_entry.data,
+    )
+
+    assert "flow_id" in result
+    flow_id = result["flow_id"]
+
+    await hass.config_entries.flow.async_configure(
+        flow_id, user_input={"next_step_id": "pick_implementation"}
+    )
+
+    # pylint: disable=protected-access
+    state = config_entry_oauth2_flow._encode_jwt(
+        hass,
+        {
+            "flow_id": flow_id,
+            "redirect_uri": "https://example.com/auth/external/callback",
+        },
+    )
+
+    client = await hass_client_no_auth()
+    await client.get(f"/auth/external/callback?code=abcd&state={state}")
+    aioclient_mock.post(
+        "https://developer.lametric.com/api/v2/oauth2/token",
+        json={
+            "refresh_token": "mock-refresh-token",
+            "access_token": "mock-access-token",
+            "type": "Bearer",
+            "expires_in": 60,
+        },
+    )
+
+    result2 = await hass.config_entries.flow.async_configure(flow_id)
+
+    assert result2.get("type") == FlowResultType.ABORT
+    assert result2.get("reason") == "reauth_successful"
+    assert mock_config_entry.data == {
+        CONF_HOST: "127.0.0.1",
+        CONF_API_KEY: "mock-api-key",
+        CONF_MAC: "AA:BB:CC:DD:EE:FF",
+    }
+
+    assert len(mock_lametric_cloud_config_flow.devices.mock_calls) == 1
+    assert len(mock_lametric_config_flow.device.mock_calls) == 1
+    assert len(mock_lametric_config_flow.notify.mock_calls) == 1
+
+
+async def test_reauth_cloud_abort_device_not_found(
+    hass: HomeAssistant,
+    hass_client_no_auth: Callable[[], Awaitable[TestClient]],
+    aioclient_mock: AiohttpClientMocker,
+    current_request_with_host: None,
+    mock_setup_entry: MagicMock,
+    mock_lametric_cloud_config_flow: MagicMock,
+    mock_lametric_config_flow: MagicMock,
+    mock_config_entry: MockConfigEntry,
+) -> None:
+    """Test reauth flow importing api keys from the cloud."""
+    mock_config_entry.unique_id = "UKNOWN_DEVICE"
+    mock_config_entry.add_to_hass(hass)
+
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        context={
+            "source": SOURCE_REAUTH,
+            "unique_id": mock_config_entry.unique_id,
+            "entry_id": mock_config_entry.entry_id,
+        },
+        data=mock_config_entry.data,
+    )
+
+    assert "flow_id" in result
+    flow_id = result["flow_id"]
+
+    await hass.config_entries.flow.async_configure(
+        flow_id, user_input={"next_step_id": "pick_implementation"}
+    )
+
+    # pylint: disable=protected-access
+    state = config_entry_oauth2_flow._encode_jwt(
+        hass,
+        {
+            "flow_id": flow_id,
+            "redirect_uri": "https://example.com/auth/external/callback",
+        },
+    )
+
+    client = await hass_client_no_auth()
+    await client.get(f"/auth/external/callback?code=abcd&state={state}")
+    aioclient_mock.post(
+        "https://developer.lametric.com/api/v2/oauth2/token",
+        json={
+            "refresh_token": "mock-refresh-token",
+            "access_token": "mock-access-token",
+            "type": "Bearer",
+            "expires_in": 60,
+        },
+    )
+
+    result2 = await hass.config_entries.flow.async_configure(flow_id)
+
+    assert result2.get("type") == FlowResultType.ABORT
+    assert result2.get("reason") == "reauth_device_not_found"
+
+    assert len(mock_lametric_cloud_config_flow.devices.mock_calls) == 1
+    assert len(mock_lametric_config_flow.device.mock_calls) == 0
+    assert len(mock_lametric_config_flow.notify.mock_calls) == 0
+
+
+async def test_reauth_manual(
+    hass: HomeAssistant,
+    mock_setup_entry: MagicMock,
+    mock_lametric_config_flow: MagicMock,
+    mock_config_entry: MockConfigEntry,
+) -> None:
+    """Test reauth flow with manual entry."""
+    mock_config_entry.add_to_hass(hass)
+
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        context={
+            "source": SOURCE_REAUTH,
+            "unique_id": mock_config_entry.unique_id,
+            "entry_id": mock_config_entry.entry_id,
+        },
+        data=mock_config_entry.data,
+    )
+
+    assert "flow_id" in result
+    flow_id = result["flow_id"]
+
+    await hass.config_entries.flow.async_configure(
+        flow_id, user_input={"next_step_id": "manual_entry"}
+    )
+
+    result2 = await hass.config_entries.flow.async_configure(
+        flow_id, user_input={CONF_API_KEY: "mock-api-key"}
+    )
+
+    assert result2.get("type") == FlowResultType.ABORT
+    assert result2.get("reason") == "reauth_successful"
+    assert mock_config_entry.data == {
+        CONF_HOST: "127.0.0.1",
+        CONF_API_KEY: "mock-api-key",
+        CONF_MAC: "AA:BB:CC:DD:EE:FF",
+    }
+
+    assert len(mock_lametric_config_flow.device.mock_calls) == 1
+    assert len(mock_lametric_config_flow.notify.mock_calls) == 1

tests/components/lametric/test_init.py

@@ -3,11 +3,15 @@
 from unittest.mock import MagicMock
 
 from aiohttp import ClientWebSocketResponse
-from demetriek import LaMetricConnectionError, LaMetricConnectionTimeoutError
+from demetriek import (
+    LaMetricAuthenticationError,
+    LaMetricConnectionError,
+    LaMetricConnectionTimeoutError,
+)
 import pytest
 
 from homeassistant.components.lametric.const import DOMAIN
-from homeassistant.config_entries import ConfigEntryState
+from homeassistant.config_entries import SOURCE_REAUTH, ConfigEntryState
 from homeassistant.const import CONF_CLIENT_ID, CONF_CLIENT_SECRET
 from homeassistant.core import HomeAssistant
 from homeassistant.setup import async_setup_component
@@ -70,3 +74,30 @@ async def test_yaml_config_raises_repairs(
     issues = await get_repairs(hass, hass_ws_client)
     assert len(issues) == 1
     assert issues[0]["issue_id"] == "manual_migration"
+
+
+async def test_config_entry_authentication_failed(
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    mock_lametric: MagicMock,
+) -> None:
+    """Test trigger reauthentication flow."""
+    mock_config_entry.add_to_hass(hass)
+
+    mock_lametric.device.side_effect = LaMetricAuthenticationError
+
+    await hass.config_entries.async_setup(mock_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert mock_config_entry.state is ConfigEntryState.SETUP_ERROR
+
+    flows = hass.config_entries.flow.async_progress()
+    assert len(flows) == 1
+
+    flow = flows[0]
+    assert flow.get("step_id") == "choice_enter_manual_or_fetch_cloud"
+    assert flow.get("handler") == DOMAIN
+
+    assert "context" in flow
+    assert flow["context"].get("source") == SOURCE_REAUTH
+    assert flow["context"].get("entry_id") == mock_config_entry.entry_id