From adfa6758d2935414a580c1c2236bce4b81bc7ee2 Mon Sep 17 00:00:00 2001
From: Indu Prakash <6459774+iprak@users.noreply.github.com>
Date: Fri, 26 Jul 2024 21:32:01 -0500
Subject: [PATCH] Remove inapplicable sensor entities (#41)

* Refactored unique id generation into coordinator
* Remove inapplicable sensor entities
---
 custom_components/weatherapi/__init__.py    | 22 ++++++++++++++++++++-
 custom_components/weatherapi/coordinator.py |  6 ++++++
 custom_components/weatherapi/sensor.py      |  2 +-
 3 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/custom_components/weatherapi/__init__.py b/custom_components/weatherapi/__init__.py
index 2c2882d..e7ebbf3 100755
--- a/custom_components/weatherapi/__init__.py
+++ b/custom_components/weatherapi/__init__.py
@@ -1,6 +1,7 @@
 """The WeatherAPI integration."""
 
 from datetime import timedelta
+import logging
 from typing import Final
 
 from homeassistant.config_entries import ConfigEntry
@@ -12,18 +13,23 @@
     Platform,
 )
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers import entity_registry as er
 
 from .const import (
+    CONFIG_ADD_SENSORS,
     CONFIG_FORECAST,
     CONFIG_IGNORE_PAST_HOUR,
+    DEFAULT_ADD_SENSORS,
     DEFAULT_FORECAST,
     DEFAULT_IGNORE_PAST_HOUR,
     DOMAIN,
     UPDATE_INTERVAL_MINUTES,
 )
 from .coordinator import WeatherAPIUpdateCoordinator, WeatherAPIUpdateCoordinatorConfig
+from .sensor import SENSOR_DESCRIPTIONS
 
 PLATFORMS: Final = [Platform.WEATHER, Platform.SENSOR]
+_LOGGER = logging.getLogger(__name__)
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry):
@@ -33,11 +39,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry):
 
     latitude = entry.data[CONF_LATITUDE]
     longitude = entry.data[CONF_LONGITUDE]
+    location_name = entry.data[CONF_NAME]
 
     config = WeatherAPIUpdateCoordinatorConfig(
         api_key=entry.data[CONF_API_KEY],
         location=f"{latitude},{longitude}",
-        name=entry.data[CONF_NAME],
+        name=location_name,
         update_interval=timedelta(minutes=UPDATE_INTERVAL_MINUTES),
         forecast=entry.options.get(CONFIG_FORECAST, DEFAULT_FORECAST),
         ignore_past_forecast=entry.options.get(
@@ -52,6 +59,19 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry):
     hass.data[DOMAIN][entry.entry_id] = coordinator
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
+    # Remove previous sensor entries if sensors are not being added
+    if not entry.options.get(CONFIG_ADD_SENSORS, DEFAULT_ADD_SENSORS):
+        ent_reg = er.async_get(hass)
+
+        for description in SENSOR_DESCRIPTIONS:
+            unique_id = coordinator.generate_sensor_unique_id(description)
+
+            if entity_id := ent_reg.async_get_entity_id(
+                Platform.SENSOR, DOMAIN, unique_id
+            ):
+                _LOGGER.debug("Removing sensor entity %s", entity_id)
+                ent_reg.async_remove(entity_id)
+
     return True
 
 
diff --git a/custom_components/weatherapi/coordinator.py b/custom_components/weatherapi/coordinator.py
index fa9d777..02be45a 100755
--- a/custom_components/weatherapi/coordinator.py
+++ b/custom_components/weatherapi/coordinator.py
@@ -21,6 +21,7 @@
     ATTR_PM_10,
     ATTR_SO2,
 )
+from homeassistant.components.sensor import SensorEntityDescription
 from homeassistant.components.weather import (
     ATTR_CONDITION_CLEAR_NIGHT,
     ATTR_CONDITION_SUNNY,
@@ -435,6 +436,11 @@ def parse_current(self, json):
             ),
         }
 
+    def generate_sensor_unique_id(self, description: SensorEntityDescription) -> str:
+        """Generate unique id for a sensor."""
+        name = f"{self.config.name} {description.name}"
+        return f"{self.location}_{name}"
+
 
 class InvalidApiKey(HomeAssistantError):
     """Error to indicate there is an invalid api key."""
diff --git a/custom_components/weatherapi/sensor.py b/custom_components/weatherapi/sensor.py
index d9b8737..cf9b0b6 100755
--- a/custom_components/weatherapi/sensor.py
+++ b/custom_components/weatherapi/sensor.py
@@ -150,7 +150,7 @@ def __init__(
             entity_id_format, f"{WEATHERAPI_DOMAIN}_{name}", hass=coordinator.hass
         )
 
-        self._attr_unique_id = f"{self.coordinator.location}_{name}"
+        self._attr_unique_id = coordinator.generate_sensor_unique_id(description)
 
     @property
     def available(self) -> bool: