Add MELCloud token refresh upon firmware upgrade (home-assistant#104391)

* Adding initial setup * Update homeassistant/components/melcloud/__init__.py Co-authored-by: G Johansson <[email protected]> * Adding ConfigEntryNotReady exception * Update homeassistant/components/melcloud/__init__.py Co-authored-by: G Johansson <[email protected]> * Update homeassistant/components/melcloud/config_flow.py Co-authored-by: G Johansson <[email protected]> * Update homeassistant/components/melcloud/__init__.py Co-authored-by: G Johansson <[email protected]> * Placing exception handling in setup_entry * Expanding test cases --------- Co-authored-by: G Johansson <[email protected]>
carasen12 · Dec 2, 2023 · 7a9c381 · 7a9c381
1 parent b48ad26
commit 7a9c381
Show file tree

Hide file tree

Showing 4 changed files with 256 additions and 15 deletions.
diff --git a/homeassistant/components/melcloud/__init__.py b/homeassistant/components/melcloud/__init__.py
@@ -6,15 +6,15 @@
 import logging
 from typing import Any
 
-from aiohttp import ClientConnectionError
+from aiohttp import ClientConnectionError, ClientResponseError
 from pymelcloud import Device, get_devices
 from pymelcloud.atw_device import Zone
 import voluptuous as vol
 
 from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry
 from homeassistant.const import CONF_TOKEN, CONF_USERNAME, Platform
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import ConfigEntryNotReady
+from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC, DeviceInfo
@@ -66,7 +66,15 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Establish connection with MELClooud."""
     conf = entry.data
-    mel_devices = await mel_devices_setup(hass, conf[CONF_TOKEN])
+    try:
+        mel_devices = await mel_devices_setup(hass, conf[CONF_TOKEN])
+    except ClientResponseError as ex:
+        if isinstance(ex, ClientResponseError) and ex.code == 401:
+            raise ConfigEntryAuthFailed from ex
+        raise ConfigEntryNotReady from ex
+    except (asyncio.TimeoutError, ClientConnectionError) as ex:
+        raise ConfigEntryNotReady from ex
+
     hass.data.setdefault(DOMAIN, {}).update({entry.entry_id: mel_devices})
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
     return True
@@ -162,17 +170,13 @@ async def mel_devices_setup(
 ) -> dict[str, list[MelCloudDevice]]:
     """Query connected devices from MELCloud."""
     session = async_get_clientsession(hass)
-    try:
-        async with asyncio.timeout(10):
-            all_devices = await get_devices(
-                token,
-                session,
-                conf_update_interval=timedelta(minutes=5),
-                device_set_debounce=timedelta(seconds=1),
-            )
-    except (asyncio.TimeoutError, ClientConnectionError) as ex:
-        raise ConfigEntryNotReady() from ex
-
+    async with asyncio.timeout(10):
+        all_devices = await get_devices(
+            token,
+            session,
+            conf_update_interval=timedelta(minutes=5),
+            device_set_debounce=timedelta(seconds=1),
+        )
     wrapped_devices: dict[str, list[MelCloudDevice]] = {}
     for device_type, devices in all_devices.items():
         wrapped_devices[device_type] = [MelCloudDevice(device) for device in devices]

diff --git a/homeassistant/components/melcloud/config_flow.py b/homeassistant/components/melcloud/config_flow.py
@@ -2,7 +2,10 @@
 from __future__ import annotations
 
 import asyncio
+from collections.abc import Mapping
 from http import HTTPStatus
+import logging
+from typing import Any
 
 from aiohttp import ClientError, ClientResponseError
 import pymelcloud
@@ -11,12 +14,14 @@
 from homeassistant import config_entries
 from homeassistant.const import CONF_PASSWORD, CONF_TOKEN, CONF_USERNAME
 from homeassistant.core import DOMAIN as HOMEASSISTANT_DOMAIN, HomeAssistant
-from homeassistant.data_entry_flow import AbortFlow, FlowResultType
+from homeassistant.data_entry_flow import AbortFlow, FlowResult, FlowResultType
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.issue_registry import IssueSeverity, async_create_issue
 
 from .const import DOMAIN
 
+_LOGGER = logging.getLogger(__name__)
+
 
 async def async_create_import_issue(
     hass: HomeAssistant, source: str, issue: str, success: bool = False
@@ -56,6 +61,8 @@ class FlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
 
     VERSION = 1
 
+    entry: config_entries.ConfigEntry | None = None
+
     async def _create_entry(self, username: str, token: str):
         """Register new entry."""
         await self.async_set_unique_id(username)
@@ -126,3 +133,67 @@ async def async_step_import(self, user_input):
         if result["type"] == FlowResultType.CREATE_ENTRY:
             await async_create_import_issue(self.hass, self.context["source"], "", True)
         return result
+
+    async def async_step_reauth(self, entry_data: Mapping[str, Any]) -> FlowResult:
+        """Handle initiation of re-authentication with MELCloud."""
+        self.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:
+        """Handle re-authentication with MELCloud."""
+        errors: dict[str, str] = {}
+
+        if user_input is not None and self.entry:
+            aquired_token, errors = await self.async_reauthenticate_client(user_input)
+
+            if not errors:
+                self.hass.config_entries.async_update_entry(
+                    self.entry,
+                    data={CONF_TOKEN: aquired_token},
+                )
+                self.hass.async_create_task(
+                    self.hass.config_entries.async_reload(self.entry.entry_id)
+                )
+                return self.async_abort(reason="reauth_successful")
+
+        return self.async_show_form(
+            step_id="reauth_confirm",
+            data_schema=vol.Schema(
+                {vol.Required(CONF_USERNAME): str, vol.Required(CONF_PASSWORD): str}
+            ),
+            errors=errors,
+        )
+
+    async def async_reauthenticate_client(
+        self, user_input: dict[str, Any]
+    ) -> tuple[str | None, dict[str, str]]:
+        """Reauthenticate with MELCloud."""
+        errors: dict[str, str] = {}
+        acquired_token = None
+
+        try:
+            async with asyncio.timeout(10):
+                acquired_token = await pymelcloud.login(
+                    user_input[CONF_USERNAME],
+                    user_input[CONF_PASSWORD],
+                    async_get_clientsession(self.hass),
+                )
+        except (ClientResponseError, AttributeError) as err:
+            if isinstance(err, ClientResponseError) and err.status in (
+                HTTPStatus.UNAUTHORIZED,
+                HTTPStatus.FORBIDDEN,
+            ):
+                errors["base"] = "invalid_auth"
+            elif isinstance(err, AttributeError) and err.name == "get":
+                errors["base"] = "invalid_auth"
+            else:
+                errors["base"] = "cannot_connect"
+        except (
+            asyncio.TimeoutError,
+            ClientError,
+        ):
+            errors["base"] = "cannot_connect"
+
+        return acquired_token, errors
diff --git a/homeassistant/components/melcloud/strings.json b/homeassistant/components/melcloud/strings.json
@@ -8,6 +8,14 @@
           "username": "[%key:common::config_flow::data::email%]",
           "password": "[%key:common::config_flow::data::password%]"
         }
+      },
+      "reauth_confirm": {
+        "title": "[%key:common::config_flow::title::reauth%]",
+        "description": "The Melcloud integration needs to re-authenticate your connection details",
+        "data": {
+          "username": "[%key:common::config_flow::data::email%]",
+          "password": "[%key:common::config_flow::data::password%]"
+        }
       }
     },
     "error": {
@@ -16,6 +24,7 @@
       "unknown": "[%key:common::config_flow::error::unknown%]"
     },
     "abort": {
+      "reauth_successful": "[%key:common::config_flow::abort::reauth_successful%]",
       "already_configured": "MELCloud integration already configured for this email. Access token has been refreshed."
     }
   },

diff --git a/tests/components/melcloud/test_config_flow.py b/tests/components/melcloud/test_config_flow.py
@@ -9,7 +9,9 @@
 
 from homeassistant import config_entries, data_entry_flow
 from homeassistant.components.melcloud.const import DOMAIN
+from homeassistant.config_entries import SOURCE_REAUTH
 from homeassistant.core import DOMAIN as HOMEASSISTANT_DOMAIN, HomeAssistant
+from homeassistant.data_entry_flow import FlowResultType
 import homeassistant.helpers.issue_registry as ir
 
 from tests.common import MockConfigEntry
@@ -287,3 +289,158 @@ async def test_token_refresh(hass: HomeAssistant, mock_login, mock_get_devices)
     entry = entries[0]
     assert entry.data["username"] == "[email protected]"
     assert entry.data["token"] == "test-token"
+
+
+async def test_token_reauthentication(
+    hass: HomeAssistant,
+    mock_login,
+    mock_get_devices,
+) -> None:
+    """Re-configuration with existing username should refresh token, if made invalid."""
+    mock_entry = MockConfigEntry(
+        domain=DOMAIN,
+        data={"username": "[email protected]", "token": "test-original-token"},
+        unique_id="[email protected]",
+    )
+    mock_entry.add_to_hass(hass)
+
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        context={
+            "source": SOURCE_REAUTH,
+            "unique_id": mock_entry.unique_id,
+            "entry_id": mock_entry.entry_id,
+        },
+        data=mock_entry.data,
+    )
+    assert result["type"] == FlowResultType.FORM
+    assert result["step_id"] == "reauth_confirm"
+
+    with patch(
+        "homeassistant.components.melcloud.async_setup_entry",
+        return_value=True,
+    ) as mock_setup_entry:
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {"username": "[email protected]", "password": "test-password"},
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] == FlowResultType.ABORT
+    assert result["reason"] == "reauth_successful"
+    assert len(mock_setup_entry.mock_calls) == 1
+
+
+@pytest.mark.parametrize(
+    ("error", "reason"),
+    [
+        (asyncio.TimeoutError(), "cannot_connect"),
+        (AttributeError(name="get"), "invalid_auth"),
+    ],
+)
+async def test_form_errors_reauthentication(
+    hass: HomeAssistant, mock_login, error, reason
+) -> None:
+    """Test we handle cannot connect error."""
+    mock_login.side_effect = error
+    mock_entry = MockConfigEntry(
+        domain=DOMAIN,
+        data={"username": "[email protected]", "token": "test-original-token"},
+        unique_id="[email protected]",
+    )
+    mock_entry.add_to_hass(hass)
+
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        context={
+            "source": SOURCE_REAUTH,
+            "unique_id": mock_entry.unique_id,
+            "entry_id": mock_entry.entry_id,
+        },
+        data=mock_entry.data,
+    )
+
+    with patch(
+        "homeassistant.components.melcloud.async_setup_entry",
+        return_value=True,
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {"username": "[email protected]", "password": "test-password"},
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] == FlowResultType.FORM
+    assert result["errors"]["base"] == reason
+
+    mock_login.side_effect = None
+    with patch(
+        "homeassistant.components.melcloud.async_setup_entry",
+        return_value=True,
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {"username": "[email protected]", "password": "test-password"},
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] == FlowResultType.ABORT
+    assert result["reason"] == "reauth_successful"
+
+
+@pytest.mark.parametrize(
+    ("error", "reason"),
+    [
+        (HTTPStatus.UNAUTHORIZED, "invalid_auth"),
+        (HTTPStatus.FORBIDDEN, "invalid_auth"),
+        (HTTPStatus.INTERNAL_SERVER_ERROR, "cannot_connect"),
+    ],
+)
+async def test_client_errors_reauthentication(
+    hass: HomeAssistant, mock_login, mock_request_info, error, reason
+) -> None:
+    """Test we handle cannot connect error."""
+    mock_login.side_effect = ClientResponseError(mock_request_info(), (), status=error)
+    mock_entry = MockConfigEntry(
+        domain=DOMAIN,
+        data={"username": "[email protected]", "token": "test-original-token"},
+        unique_id="[email protected]",
+    )
+    mock_entry.add_to_hass(hass)
+
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        context={
+            "source": SOURCE_REAUTH,
+            "unique_id": mock_entry.unique_id,
+            "entry_id": mock_entry.entry_id,
+        },
+        data=mock_entry.data,
+    )
+
+    with patch(
+        "homeassistant.components.melcloud.async_setup_entry",
+        return_value=True,
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {"username": "[email protected]", "password": "test-password"},
+        )
+        await hass.async_block_till_done()
+
+    assert result["errors"]["base"] == reason
+    assert result["type"] == FlowResultType.FORM
+
+    mock_login.side_effect = None
+    with patch(
+        "homeassistant.components.melcloud.async_setup_entry",
+        return_value=True,
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {"username": "[email protected]", "password": "test-password"},
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] == FlowResultType.ABORT
+    assert result["reason"] == "reauth_successful"