From bdad37f0caddef10a3e34c9e2c4094413aeffde6 Mon Sep 17 00:00:00 2001
From: James Smith <james@loopj.com>
Date: Wed, 8 Jan 2025 13:18:55 -0800
Subject: [PATCH] Remove BackBox objects from device registry, special-case
 parent relationship

---
 custom_components/vantage/__init__.py |  4 ++++
 custom_components/vantage/device.py   |  7 +++++--
 custom_components/vantage/migrate.py  | 30 +++++++++++++++++++++++++++
 3 files changed, 39 insertions(+), 2 deletions(-)
 create mode 100644 custom_components/vantage/migrate.py

diff --git a/custom_components/vantage/__init__.py b/custom_components/vantage/__init__.py
index 573dae8..3b809b5 100644
--- a/custom_components/vantage/__init__.py
+++ b/custom_components/vantage/__init__.py
@@ -30,6 +30,7 @@
 from .device import async_setup_devices
 from .entity import async_cleanup_entities
 from .events import async_setup_events
+from .migrate import async_migrate_data
 from .services import async_register_services
 
 PLATFORMS: list[Platform] = [
@@ -110,6 +111,9 @@ async def handle_system_program_event(
         # retrying set up later.
         raise ConfigEntryNotReady from err
 
+    # Run any migrations
+    await async_migrate_data(hass, entry)
+
     return True
 
 
diff --git a/custom_components/vantage/device.py b/custom_components/vantage/device.py
index e75b67d..d413016 100644
--- a/custom_components/vantage/device.py
+++ b/custom_components/vantage/device.py
@@ -52,7 +52,6 @@ def handle_device_event(event_type: VantageEvent, obj: T, _data: Any) -> None:
         entry.async_on_unload(controller.subscribe(handle_device_event))
 
     # Register "parent" devices (controllers, modules, port devices, and stations)
-    register_items(vantage.back_boxes)
     register_items(vantage.masters)
     register_items(vantage.modules)
     register_items(vantage.port_devices)
@@ -102,7 +101,11 @@ def vantage_device_info(client: Vantage, obj: SystemObject) -> DeviceInfo:
 
     # Set up device relationships
     if not isinstance(obj, Master):
-        if isinstance(obj, ChildObject) and obj.parent.id in client:
+        if (
+            isinstance(obj, ChildObject)
+            and obj.parent.id in client
+            and not client.back_boxes.get(obj.parent.id)
+        ):
             device_info["via_device"] = (DOMAIN, str(obj.parent.id))
         else:
             device_info["via_device"] = (DOMAIN, str(obj.master_id))
diff --git a/custom_components/vantage/migrate.py b/custom_components/vantage/migrate.py
new file mode 100644
index 0000000..dbf79b6
--- /dev/null
+++ b/custom_components/vantage/migrate.py
@@ -0,0 +1,30 @@
+"""Migration functions for the Vantage integration."""
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import device_registry as dr
+
+from .const import LOGGER
+
+
+async def async_migrate_data(hass: HomeAssistant, entry: ConfigEntry) -> None:
+    """Run all Vantage data migrations."""
+
+    async_delete_back_boxes(hass, entry)
+
+
+def async_delete_back_boxes(hass: HomeAssistant, entry: ConfigEntry) -> None:
+    """Delete back boxes from the device registry."""
+    dev_reg = dr.async_get(hass)
+
+    back_box_devices = [
+        device
+        for device in dr.async_entries_for_config_entry(dev_reg, entry.entry_id)
+        if device.model == "BackBox"
+    ]
+
+    if back_box_devices:
+        LOGGER.debug(f"Deleting {len(back_box_devices)} back boxes from the registry.")
+
+        for device in back_box_devices:
+            dev_reg.async_remove_device(device.id)