From 15ed329108af095e21f1d8eb987430a7a6e707d4 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sat, 25 Jun 2022 11:02:38 -0500
Subject: [PATCH] Add async_remove_config_entry_device support to nexia
 (#73966)

Co-authored-by: Paulus Schoutsen <balloob@gmail.com>
---
 homeassistant/components/nexia/__init__.py | 23 +++++++--
 tests/components/nexia/test_init.py        | 60 ++++++++++++++++++++++
 2 files changed, 80 insertions(+), 3 deletions(-)
 create mode 100644 tests/components/nexia/test_init.py

diff --git a/homeassistant/components/nexia/__init__.py b/homeassistant/components/nexia/__init__.py
index ab491c0a271..355c17a2ed1 100644
--- a/homeassistant/components/nexia/__init__.py
+++ b/homeassistant/components/nexia/__init__.py
@@ -5,11 +5,13 @@ import logging
 import aiohttp
 from nexia.const import BRAND_NEXIA
 from nexia.home import NexiaHome
+from nexia.thermostat import NexiaThermostat
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
+from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 import homeassistant.helpers.config_validation as cv
 
@@ -66,8 +68,23 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-    if unload_ok:
+    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
         hass.data[DOMAIN].pop(entry.entry_id)
-
     return unload_ok
+
+
+async def async_remove_config_entry_device(
+    hass: HomeAssistant, config_entry: ConfigEntry, device_entry: dr.DeviceEntry
+) -> bool:
+    """Remove a nexia config entry from a device."""
+    coordinator: NexiaDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id]
+    nexia_home: NexiaHome = coordinator.nexia_home
+    dev_ids = {dev_id[1] for dev_id in device_entry.identifiers if dev_id[0] == DOMAIN}
+    for thermostat_id in nexia_home.get_thermostat_ids():
+        if thermostat_id in dev_ids:
+            return False
+        thermostat: NexiaThermostat = nexia_home.get_thermostat_by_id(thermostat_id)
+        for zone_id in thermostat.get_zone_ids():
+            if zone_id in dev_ids:
+                return False
+    return True
diff --git a/tests/components/nexia/test_init.py b/tests/components/nexia/test_init.py
new file mode 100644
index 00000000000..667c03a23cf
--- /dev/null
+++ b/tests/components/nexia/test_init.py
@@ -0,0 +1,60 @@
+"""The init tests for the nexia platform."""
+
+
+from homeassistant.components.nexia.const import DOMAIN
+from homeassistant.helpers import device_registry as dr, entity_registry as er
+from homeassistant.helpers.entity_registry import EntityRegistry
+from homeassistant.setup import async_setup_component
+
+from .util import async_init_integration
+
+
+async def remove_device(ws_client, device_id, config_entry_id):
+    """Remove config entry from a device."""
+    await ws_client.send_json(
+        {
+            "id": 5,
+            "type": "config/device_registry/remove_config_entry",
+            "config_entry_id": config_entry_id,
+            "device_id": device_id,
+        }
+    )
+    response = await ws_client.receive_json()
+    return response["success"]
+
+
+async def test_device_remove_devices(hass, hass_ws_client):
+    """Test we can only remove a device that no longer exists."""
+    await async_setup_component(hass, "config", {})
+    config_entry = await async_init_integration(hass)
+    entry_id = config_entry.entry_id
+    device_registry = dr.async_get(hass)
+
+    registry: EntityRegistry = er.async_get(hass)
+    entity = registry.entities["sensor.nick_office_temperature"]
+
+    live_zone_device_entry = device_registry.async_get(entity.device_id)
+    assert (
+        await remove_device(
+            await hass_ws_client(hass), live_zone_device_entry.id, entry_id
+        )
+        is False
+    )
+
+    entity = registry.entities["sensor.master_suite_relative_humidity"]
+    live_thermostat_device_entry = device_registry.async_get(entity.device_id)
+    assert (
+        await remove_device(
+            await hass_ws_client(hass), live_thermostat_device_entry.id, entry_id
+        )
+        is False
+    )
+
+    dead_device_entry = device_registry.async_get_or_create(
+        config_entry_id=config_entry.entry_id,
+        identifiers={(DOMAIN, "unused")},
+    )
+    assert (
+        await remove_device(await hass_ws_client(hass), dead_device_entry.id, entry_id)
+        is True
+    )