Add unique_id to senz config_entry (home-assistant#156472)

Author: astrandb

homeassistant/components/senz/__init__.py

@@ -7,6 +7,7 @@
 
 from aiosenz import SENZAPI, Thermostat
 from httpx import RequestError
+import jwt
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
@@ -82,3 +83,27 @@ async def update_thermostats() -> dict[str, Thermostat]:
 async def async_unload_entry(hass: HomeAssistant, entry: SENZConfigEntry) -> bool:
     """Unload a config entry."""
     return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
+
+
+async def async_migrate_entry(
+    hass: HomeAssistant, config_entry: SENZConfigEntry
+) -> bool:
+    """Migrate old entry."""
+
+    # Use sub(ject) from access_token as unique_id
+    if config_entry.version == 1 and config_entry.minor_version == 1:
+        token = jwt.decode(
+            config_entry.data["token"]["access_token"],
+            options={"verify_signature": False},
+        )
+        uid = token["sub"]
+        hass.config_entries.async_update_entry(
+            config_entry, unique_id=uid, minor_version=2
+        )
+        _LOGGER.info(
+            "Migration to version %s.%s successful",
+            config_entry.version,
+            config_entry.minor_version,
+        )
+
+    return True

homeassistant/components/senz/config_flow.py

@@ -2,6 +2,9 @@
 
 import logging
 
+import jwt
+
+from homeassistant.config_entries import ConfigFlowResult
 from homeassistant.helpers import config_entry_oauth2_flow
 
 from .const import DOMAIN
@@ -12,6 +15,8 @@ class OAuth2FlowHandler(
 ):
     """Config flow to handle SENZ OAuth2 authentication."""
 
+    VERSION = 1
+    MINOR_VERSION = 2
     DOMAIN = DOMAIN
 
     @property
@@ -23,3 +28,15 @@ def logger(self) -> logging.Logger:
     def extra_authorize_data(self) -> dict:
         """Extra data that needs to be appended to the authorize url."""
         return {"scope": "restapi offline_access"}
+
+    async def async_oauth_create_entry(self, data: dict) -> ConfigFlowResult:
+        """Create or update the config entry."""
+
+        token = jwt.decode(
+            data["token"]["access_token"], options={"verify_signature": False}
+        )
+        uid = token["sub"]
+        await self.async_set_unique_id(uid)
+
+        self._abort_if_unique_id_configured()
+        return await super().async_oauth_create_entry(data)

tests/components/senz/conftest.py

@@ -14,9 +14,10 @@
 )
 from homeassistant.components.senz.const import DOMAIN
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers import config_entry_oauth2_flow
 from homeassistant.setup import async_setup_component
 
-from .const import CLIENT_ID, CLIENT_SECRET
+from .const import CLIENT_ID, CLIENT_SECRET, ENTRY_UNIQUE_ID
 
 from tests.common import (
     MockConfigEntry,
@@ -63,7 +64,7 @@ def mock_expires_at() -> float:
 def mock_config_entry(hass: HomeAssistant, expires_at: float) -> MockConfigEntry:
     """Return the default mocked config entry."""
     config_entry = MockConfigEntry(
-        minor_version=1,
+        minor_version=2,
         domain=DOMAIN,
         title="Senz test",
         data={
@@ -77,6 +78,7 @@ def mock_config_entry(hass: HomeAssistant, expires_at: float) -> MockConfigEntry
             },
         },
         entry_id="senz_test",
+        unique_id=ENTRY_UNIQUE_ID,
     )
     config_entry.add_to_hass(hass)
     return config_entry
@@ -109,3 +111,20 @@ async def setup_credentials(hass: HomeAssistant) -> None:
         ),
         DOMAIN,
     )
+
+
+@pytest.fixture
+async def access_token(hass: HomeAssistant) -> str:
+    """Return a valid access token."""
+    return config_entry_oauth2_flow._encode_jwt(
+        hass,
+        {
+            "sub": ENTRY_UNIQUE_ID,
+            "aud": [],
+            "scp": [
+                "rest_api",
+                "offline_access",
+            ],
+            "ou_code": "NA",
+        },
+    )

tests/components/senz/const.py

@@ -2,3 +2,5 @@
 
 CLIENT_ID = "test_client_id"
 CLIENT_SECRET = "test_client_secret"
+
+ENTRY_UNIQUE_ID = "test_unique_id"

tests/components/senz/test_config_flow.py

@@ -12,11 +12,13 @@
 )
 from homeassistant.components.senz.const import DOMAIN
 from homeassistant.core import HomeAssistant
+from homeassistant.data_entry_flow import FlowResultType
 from homeassistant.helpers import config_entry_oauth2_flow
 from homeassistant.setup import async_setup_component
 
 from .const import CLIENT_ID, CLIENT_SECRET
 
+from tests.common import MockConfigEntry
 from tests.test_util.aiohttp import AiohttpClientMocker
 from tests.typing import ClientSessionGenerator
 
@@ -26,6 +28,7 @@ async def test_full_flow(
     hass: HomeAssistant,
     hass_client_no_auth: ClientSessionGenerator,
     aioclient_mock: AiohttpClientMocker,
+    access_token: str,
 ) -> None:
     """Check full flow."""
     await async_setup_component(hass, DOMAIN, {})
@@ -61,7 +64,7 @@ async def test_full_flow(
         TOKEN_ENDPOINT,
         json={
             "refresh_token": "mock-refresh-token",
-            "access_token": "mock-access-token",
+            "access_token": access_token,
             "type": "Bearer",
             "expires_in": 60,
         },
@@ -74,3 +77,52 @@ async def test_full_flow(
 
     assert len(hass.config_entries.async_entries(DOMAIN)) == 1
     assert len(mock_setup.mock_calls) == 1
+
+
+@pytest.mark.usefixtures("current_request_with_host")
+async def test_duplicate_flow(
+    hass: HomeAssistant,
+    hass_client_no_auth: ClientSessionGenerator,
+    aioclient_mock: AiohttpClientMocker,
+    mock_config_entry: MockConfigEntry,
+    access_token: str,
+) -> None:
+    """Check full flow with duplicate entry."""
+    mock_config_entry.add_to_hass(hass)
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+    state = config_entry_oauth2_flow._encode_jwt(
+        hass,
+        {
+            "flow_id": result["flow_id"],
+            "redirect_uri": "https://example.com/auth/external/callback",
+        },
+    )
+
+    assert result["url"] == (
+        f"{AUTHORIZATION_ENDPOINT}?response_type=code&client_id={CLIENT_ID}"
+        "&redirect_uri=https://example.com/auth/external/callback"
+        f"&state={state}&scope=restapi+offline_access"
+    )
+
+    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(
+        TOKEN_ENDPOINT,
+        json={
+            "refresh_token": "mock-refresh-token",
+            "access_token": access_token,
+            "type": "Bearer",
+            "expires_in": 60,
+        },
+    )
+
+    with patch("homeassistant.components.senz.async_setup_entry", return_value=True):
+        result2 = await hass.config_entries.flow.async_configure(result["flow_id"])
+
+    assert result2["type"] is FlowResultType.ABORT
+    assert result2["reason"] == "already_configured"

tests/components/senz/test_init.py

@@ -2,13 +2,15 @@
 
 from unittest.mock import MagicMock, patch
 
+from homeassistant.components.senz.const import DOMAIN
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.config_entry_oauth2_flow import (
     ImplementationUnavailableError,
 )
 
 from . import setup_integration
+from .const import ENTRY_UNIQUE_ID
 
 from tests.common import MockConfigEntry
 
@@ -43,3 +45,36 @@ async def test_oauth_implementation_not_available(
         await hass.async_block_till_done()
 
     assert mock_config_entry.state is ConfigEntryState.SETUP_RETRY
+
+
+async def test_migrate_config_entry(
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    mock_senz_client: MagicMock,
+    expires_at: float,
+    access_token: str,
+) -> None:
+    """Test migration of config entry."""
+    mock_entry_v1_1 = MockConfigEntry(
+        version=1,
+        minor_version=1,
+        domain=DOMAIN,
+        title="SENZ test",
+        data={
+            "auth_implementation": DOMAIN,
+            "token": {
+                "access_token": access_token,
+                "scope": "rest_api offline_access",
+                "expires_in": 86399,
+                "refresh_token": "3012bc9f-7a65-4240-b817-9154ffdcc30f",
+                "token_type": "Bearer",
+                "expires_at": expires_at,
+            },
+        },
+        entry_id="senz_test",
+    )
+
+    await setup_integration(hass, mock_entry_v1_1)
+    assert mock_entry_v1_1.version == 1
+    assert mock_entry_v1_1.minor_version == 2
+    assert mock_entry_v1_1.unique_id == ENTRY_UNIQUE_ID