From 1b8bda6067cb73eda9e9ad312874e1e15c4d64f4 Mon Sep 17 00:00:00 2001
From: Richard Kroegel <42204099+rikroe@users.noreply.github.com>
Date: Mon, 12 Feb 2024 10:52:08 +0100
Subject: [PATCH] Remove old BMW vehicles/devices automatically (#110255)

* Remove not assigned vehicles from DeviceRegistry on startup

* Replace async_remove_device with async_update_device

* Add test

* Use generator

---------

Co-authored-by: Richard <rikroe@users.noreply.github.com>
Co-authored-by: Martin Hjelmare <marhje52@gmail.com>
---
 .../bmw_connected_drive/__init__.py           | 18 ++++++++-
 .../bmw_connected_drive/test_init.py          | 39 ++++++++++++++++++-
 2 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/homeassistant/components/bmw_connected_drive/__init__.py b/homeassistant/components/bmw_connected_drive/__init__.py
index d5a213256c3..079563b1ad3 100644
--- a/homeassistant/components/bmw_connected_drive/__init__.py
+++ b/homeassistant/components/bmw_connected_drive/__init__.py
@@ -10,7 +10,11 @@ import voluptuous as vol
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_DEVICE_ID, CONF_ENTITY_ID, CONF_NAME, Platform
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers import discovery, entity_registry as er
+from homeassistant.helpers import (
+    device_registry as dr,
+    discovery,
+    entity_registry as er,
+)
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.typing import ConfigType
@@ -146,6 +150,18 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         )
     )
 
+    # Clean up vehicles which are not assigned to the account anymore
+    account_vehicles = {(DOMAIN, v.vin) for v in coordinator.account.vehicles}
+    device_registry = dr.async_get(hass)
+    device_entries = dr.async_entries_for_config_entry(
+        device_registry, config_entry_id=entry.entry_id
+    )
+    for device in device_entries:
+        if not device.identifiers.intersection(account_vehicles):
+            device_registry.async_update_device(
+                device.id, remove_config_entry_id=entry.entry_id
+            )
+
     return True
 
 
diff --git a/tests/components/bmw_connected_drive/test_init.py b/tests/components/bmw_connected_drive/test_init.py
index bc02437f5ba..c57bfc5a9b0 100644
--- a/tests/components/bmw_connected_drive/test_init.py
+++ b/tests/components/bmw_connected_drive/test_init.py
@@ -2,11 +2,12 @@
 from unittest.mock import patch
 
 import pytest
+import respx
 
 from homeassistant.components.bmw_connected_drive.const import DOMAIN as BMW_DOMAIN
 from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry as er
+from homeassistant.helpers import device_registry as dr, entity_registry as er
 
 from . import FIXTURE_CONFIG_ENTRY
 
@@ -133,3 +134,39 @@ async def test_dont_migrate_unique_ids(
     assert entity_not_changed.unique_id == new_unique_id
 
     assert entity_migrated != entity_not_changed
+
+
+async def test_remove_stale_devices(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    bmw_fixture: respx.Router,
+) -> None:
+    """Test remove stale device registry entries."""
+    mock_config_entry = MockConfigEntry(**FIXTURE_CONFIG_ENTRY)
+    mock_config_entry.add_to_hass(hass)
+
+    device_registry.async_get_or_create(
+        config_entry_id=mock_config_entry.entry_id,
+        identifiers={(BMW_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 == {(BMW_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
+    )
+
+    # Check that the test vehicles are still available but not the stale device
+    assert len(device_entries) > 0
+    remaining_device_identifiers = set().union(*(d.identifiers for d in device_entries))
+    assert not {(BMW_DOMAIN, "stale_device_id")}.intersection(
+        remaining_device_identifiers
+    )