Clean up YoLink entities on startup (home-assistant#148718)

Author: matrixd2

homeassistant/components/yolink/__init__.py

@@ -165,6 +165,20 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     hass.data[DOMAIN][entry.entry_id] = YoLinkHomeStore(
         yolink_home, device_coordinators
     )
+
+    # Clean up yolink devices which are not associated to the account anymore.
+    device_registry = dr.async_get(hass)
+    device_entries = dr.async_entries_for_config_entry(device_registry, entry.entry_id)
+    for device_entry in device_entries:
+        for identifier in device_entry.identifiers:
+            if (
+                identifier[0] == DOMAIN
+                and device_coordinators.get(identifier[1]) is None
+            ):
+                device_registry.async_update_device(
+                    device_entry.id, remove_config_entry_id=entry.entry_id
+                )
+
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     async def async_yolink_unload(event) -> None:

tests/components/yolink/conftest.py

@@ -0,0 +1,77 @@
+"""Provide common fixtures for the YoLink integration tests."""
+
+from __future__ import annotations
+
+from collections.abc import Generator
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from yolink.home_manager import YoLinkHome
+
+from homeassistant.components.application_credentials import (
+    ClientCredential,
+    async_import_client_credential,
+)
+from homeassistant.components.yolink.api import ConfigEntryAuth
+from homeassistant.core import HomeAssistant
+from homeassistant.setup import async_setup_component
+
+from tests.common import MockConfigEntry
+
+CLIENT_ID = "12345"
+CLIENT_SECRET = "6789"
+DOMAIN = "yolink"
+
+
+@pytest.fixture
+async def setup_credentials(hass: HomeAssistant) -> None:
+    """Fixture to setup credentials."""
+    assert await async_setup_component(hass, "application_credentials", {})
+    await async_import_client_credential(
+        hass,
+        DOMAIN,
+        ClientCredential(CLIENT_ID, CLIENT_SECRET),
+    )
+
+
+@pytest.fixture(name="mock_auth_manager")
+def mock_auth_manager() -> Generator[MagicMock]:
+    """Mock the authentication manager."""
+    with patch(
+        "homeassistant.components.yolink.api.ConfigEntryAuth", autospec=True
+    ) as mock_auth:
+        mock_auth.return_value = MagicMock(spec=ConfigEntryAuth)
+        yield mock_auth
+
+
+@pytest.fixture(name="mock_yolink_home")
+def mock_yolink_home() -> Generator[AsyncMock]:
+    """Mock YoLink home instance."""
+    with patch(
+        "homeassistant.components.yolink.YoLinkHome", autospec=True
+    ) as mock_home:
+        mock_home.return_value = AsyncMock(spec=YoLinkHome)
+        yield mock_home
+
+
+@pytest.fixture
+def mock_config_entry(hass: HomeAssistant) -> MockConfigEntry:
+    """Mock a config entry for YoLink."""
+    config_entry = MockConfigEntry(
+        unique_id=DOMAIN,
+        domain=DOMAIN,
+        title="yolink",
+        data={
+            "auth_implementation": DOMAIN,
+            "token": {
+                "refresh_token": "mock-refresh-token",
+                "access_token": "mock-access-token",
+                "type": "Bearer",
+                "expires_in": 60,
+                "scope": "create",
+            },
+        },
+        options={},
+    )
+    config_entry.add_to_hass(hass)
+    return config_entry

tests/components/yolink/test_init.py

@@ -0,0 +1,38 @@
+"""Tests for the yolink integration."""
+
+import pytest
+
+from homeassistant.components.yolink import DOMAIN
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import device_registry as dr
+
+from tests.common import MockConfigEntry
+
+
+@pytest.mark.usefixtures("setup_credentials", "mock_auth_manager", "mock_yolink_home")
+async def test_device_remove_devices(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    mock_config_entry: MockConfigEntry,
+) -> None:
+    """Test we can only remove a device that no longer exists."""
+
+    device_registry.async_get_or_create(
+        config_entry_id=mock_config_entry.entry_id,
+        identifiers={(DOMAIN, "stale_device_id")},
+    )
+    device_entries = dr.async_entries_for_config_entry(
+        device_registry, mock_config_entry.entry_id
+    )
+
+    assert len(device_entries) == 1
+    device_entry = device_entries[0]
+    assert device_entry.identifiers == {(DOMAIN, "stale_device_id")}
+
+    assert await hass.config_entries.async_setup(mock_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    device_entries = dr.async_entries_for_config_entry(
+        device_registry, mock_config_entry.entry_id
+    )
+    assert len(device_entries) == 0