From f684d96e8f1aa68cc20444fb6e1b23926d33d1bd Mon Sep 17 00:00:00 2001
From: James Smith <james@loopj.com>
Date: Wed, 8 Jan 2025 13:59:23 -0800
Subject: [PATCH] Populate serial_number for Master and StationObject devices,
 remove serial number entity

---
 custom_components/vantage/device.py  | 13 ++++++++++++-
 custom_components/vantage/migrate.py | 23 +++++++++++++++++++++++
 custom_components/vantage/sensor.py  | 18 ------------------
 3 files changed, 35 insertions(+), 19 deletions(-)

diff --git a/custom_components/vantage/device.py b/custom_components/vantage/device.py
index d413016..facffb1 100644
--- a/custom_components/vantage/device.py
+++ b/custom_components/vantage/device.py
@@ -4,7 +4,13 @@
 
 from aiovantage import Vantage, VantageEvent
 from aiovantage.controllers import BaseController
-from aiovantage.models import LocationObject, Master, Parent, SystemObject
+from aiovantage.models import (
+    LocationObject,
+    Master,
+    Parent,
+    StationObject,
+    SystemObject,
+)
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
@@ -99,6 +105,11 @@ def vantage_device_info(client: Vantage, obj: SystemObject) -> DeviceInfo:
     ):
         device_info["suggested_area"] = area.name
 
+    # Attach serial number for devices that have one
+    if isinstance(obj, Master) or isinstance(obj, StationObject):
+        if obj.serial_number:
+            device_info["serial_number"] = str(obj.serial_number)
+
     # Set up device relationships
     if not isinstance(obj, Master):
         if (
diff --git a/custom_components/vantage/migrate.py b/custom_components/vantage/migrate.py
index dbf79b6..43dd545 100644
--- a/custom_components/vantage/migrate.py
+++ b/custom_components/vantage/migrate.py
@@ -3,6 +3,7 @@
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr
+from homeassistant.helpers import entity_registry as er
 
 from .const import LOGGER
 
@@ -11,6 +12,7 @@ async def async_migrate_data(hass: HomeAssistant, entry: ConfigEntry) -> None:
     """Run all Vantage data migrations."""
 
     async_delete_back_boxes(hass, entry)
+    async_delete_serial_number_entities(hass, entry)
 
 
 def async_delete_back_boxes(hass: HomeAssistant, entry: ConfigEntry) -> None:
@@ -28,3 +30,24 @@ def async_delete_back_boxes(hass: HomeAssistant, entry: ConfigEntry) -> None:
 
         for device in back_box_devices:
             dev_reg.async_remove_device(device.id)
+
+
+def async_delete_serial_number_entities(
+    hass: HomeAssistant, entry: ConfigEntry
+) -> None:
+    """Delete serial number entities from the entity registry."""
+    ent_reg = er.async_get(hass)
+
+    serial_number_entities = [
+        entity
+        for entity in er.async_entries_for_config_entry(ent_reg, entry.entry_id)
+        if entity.unique_id.endswith(":serial_number")
+    ]
+
+    if serial_number_entities:
+        LOGGER.debug(
+            f"Deleting {len(serial_number_entities)} serial number entities from the registry."
+        )
+
+        for entity in serial_number_entities:
+            ent_reg.async_remove(entity.entity_id)
diff --git a/custom_components/vantage/sensor.py b/custom_components/vantage/sensor.py
index ed07568..972a5cb 100644
--- a/custom_components/vantage/sensor.py
+++ b/custom_components/vantage/sensor.py
@@ -41,7 +41,6 @@ async def async_setup_entry(
     register_items(vantage.anemo_sensors, VantageWindSensor)
     register_items(vantage.light_sensors, VantageLightSensor)
     register_items(vantage.omni_sensors, VantageOmniSensor)
-    register_items(vantage.masters, VantageMasterSerial)
     register_items(vantage.masters, VantageMasterIP)
 
 
@@ -130,23 +129,6 @@ def native_value(self) -> int | Decimal | None:
         return self.obj.level
 
 
-class VantageMasterSerial(VantageEntity[Master], SensorEntity):
-    """Vantage controller serial number sensor entity."""
-
-    _attr_icon = "mdi:barcode"
-    _attr_entity_category = EntityCategory.DIAGNOSTIC
-
-    @property
-    def name(self) -> str | None:
-        """Return the name of the entity."""
-        return "Serial Number"
-
-    def __post_init__(self) -> None:
-        """Initialize a Vantage master serial number."""
-        self._attr_unique_id = f"{self.obj.id}:serial_number"
-        self._attr_native_value = str(self.obj.serial_number)
-
-
 class VantageMasterIP(VantageEntity[Master], SensorEntity):
     """Vantage controller IP address sensor entity."""