Merge pull request #226 from Snuffy2/Optimization-final-cleanup

Optimize entities final cleanup

custom_components/opnsense/__init__.py

@@ -1,9 +1,9 @@
 """Support for OPNsense."""
 
+import logging
 from collections.abc import Mapping
 from datetime import timedelta
-import logging
-from typing import Any, Callable
+from typing import Any
 
 import awesomeversion
 from homeassistant.config_entries import ConfigEntry
@@ -14,11 +14,9 @@
     CONF_USERNAME,
     CONF_VERIFY_SSL,
 )
-from homeassistant.core import HomeAssistant, callback
+from homeassistant.core import HomeAssistant
 from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers.aiohttp_client import async_create_clientsession
-from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.helpers.entity_registry import async_get
 from homeassistant.helpers.issue_registry import IssueSeverity, async_create_issue
 from homeassistant.helpers.typing import ConfigType
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
@@ -203,83 +201,35 @@ async def async_migrate_entry(hass: HomeAssistant, config_entry: ConfigEntry) ->
     return True
 
 
-class CoordinatorEntityManager:
-
-    def __init__(
-        self,
-        hass: HomeAssistant,
-        coordinator: OPNsenseDataUpdateCoordinator,
-        config_entry: ConfigEntry,
-        process_entities_callback: Callable,
-        async_add_entities: AddEntitiesCallback,
-    ) -> None:
-        self.hass = hass
-        self.coordinator: OPNsenseDataUpdateCoordinator = coordinator
-        self.config_entry = config_entry
-        self.process_entities_callback = process_entities_callback
-        self.async_add_entities = async_add_entities
-        hass.data[DOMAIN][config_entry.entry_id][UNDO_UPDATE_LISTENER].append(
-            coordinator.async_add_listener(self.process_entities)
-        )
-        self.entity_unique_ids = set()
-        self.entities = {}
-
-    @callback
-    def process_entities(self):
-        entities = self.process_entities_callback(self.hass, self.config_entry)
-        i_entity_unique_ids = set()
-        for entity in entities:
-            unique_id = entity.unique_id
-            if unique_id is None:
-                raise ValueError("unique_id is missing from entity")
-            i_entity_unique_ids.add(unique_id)
-            if unique_id not in self.entity_unique_ids:
-                self.async_add_entities([entity])
-                self.entity_unique_ids.add(unique_id)
-                self.entities[unique_id] = entity
-                # print(f"{unique_id} registered")
-            else:
-                # print(f"{unique_id} already registered")
-                pass
-
-        # check for missing entities
-        for entity_unique_id in self.entity_unique_ids:
-            if entity_unique_id not in i_entity_unique_ids:
-                pass
-                # print("should remove entity: " + str(self.entities[entity_unique_id].entry_id))
-                # print("candidate to remove entity: " + str(entity_unique_id))
-                # self.async_remove_entity(self.entities[entity_unique_id])
-                # self.entity_unique_ids.remove(entity_unique_id)
-                # del self.entities[entity_unique_id]
-
-    async def async_remove_entity(self, entity):
-        registry = await async_get(self.hass)
-        if entity.entity_id in registry.entities:
-            registry.async_remove(entity.entity_id)
-
-
 class OPNsenseEntity(CoordinatorEntity[OPNsenseDataUpdateCoordinator]):
     """Base entity for OPNsense"""
 
     def __init__(
         self,
         config_entry: ConfigEntry,
         coordinator: OPNsenseDataUpdateCoordinator,
-        unique_id_suffix: str,
+        unique_id_suffix: str | None = None,
         name_suffix: str | None = None,
     ) -> None:
         self.config_entry: ConfigEntry = config_entry
         self.coordinator: OPNsenseDataUpdateCoordinator = coordinator
-        self._attr_unique_id: str = slugify(
-            f"{self.opnsense_device_unique_id}_{unique_id_suffix}"
-        )
+        if unique_id_suffix:
+            self._attr_unique_id: str = slugify(
+                f"{self.opnsense_device_unique_id}_{unique_id_suffix}"
+            )
         if name_suffix:
             self._attr_name: str = (
                 f"{self.opnsense_device_name or 'OPNsense'} {name_suffix}"
             )
         self._client: OPNsenseClient | None = None
+        self._attr_extra_state_attributes: Mapping[str, Any] = {}
+        self._available: bool = False
         super().__init__(self.coordinator, self._attr_unique_id)
 
+    @property
+    def available(self) -> bool:
+        return self._available
+
     @property
     def device_info(self) -> Mapping[str, Any]:
         """Device info for the firewall."""
@@ -330,3 +280,4 @@ async def async_added_to_hass(self) -> None:
             self._client: OPNsenseClient = self._get_opnsense_client()
         if self._client is None:
             _LOGGER.error("Unable to get client in async_added_to_hass.")
+        self._handle_coordinator_update()

custom_components/opnsense/binary_sensor.py

@@ -1,8 +1,6 @@
 """OPNsense integration."""
 
-from collections.abc import Mapping
 import logging
-from typing import Any
 
 from homeassistant.components.binary_sensor import (
     BinarySensorDeviceClass,
@@ -80,20 +78,6 @@ def __init__(
         )
         self.entity_description: BinarySensorEntityDescription = entity_description
         self._attr_is_on: bool = False
-        self._attr_extra_state_attributes: Mapping[str, Any] = {}
-        self._available: bool = (
-            False  # Move this to OPNsenseEntity once all entity-types are updated
-        )
-
-    # Move this to OPNsenseEntity once all entity-types are updated
-    @property
-    def available(self) -> bool:
-        return self._available
-
-    # Move this to OPNsenseEntity once all entity-types are updated
-    async def async_added_to_hass(self) -> None:
-        await super().async_added_to_hass()
-        self._handle_coordinator_update()
 
 
 class OPNsenseCarpStatusBinarySensor(OPNsenseBinarySensor):

custom_components/opnsense/device_tracker.py

@@ -1,8 +1,8 @@
 """Support for tracking for OPNsense devices."""
 
-from datetime import datetime, timedelta
 import logging
 import time
+from datetime import datetime, timedelta
 from typing import Any, Mapping
 
 from homeassistant.components.device_tracker import SourceType
@@ -12,6 +12,8 @@
 from homeassistant.helpers import entity_platform
 from homeassistant.helpers.device_registry import (
     CONNECTION_NETWORK_MAC,
+)
+from homeassistant.helpers.device_registry import (
     async_get as async_get_dev_reg,
 )
 from homeassistant.helpers.entity import DeviceInfo
@@ -135,27 +137,18 @@ def __init__(
         hostname: str | None,
     ) -> None:
         """Set up the OPNsense scanner entity."""
-        super().__init__(config_entry, coordinator, unique_id_suffix=f"mac_{mac}")
+        super().__init__(config_entry, coordinator)
         self._mac_vendor: str | None = mac_vendor
         self._attr_name: str | None = f"{self.opnsense_device_name} {hostname or mac}"
         self._last_known_ip: str | None = None
         self._last_known_hostname: str | None = None
         self._is_connected: bool = False
         self._last_known_connected_time: str | None = None
         self._attr_entity_registry_enabled_default: bool = enabled_default
-        self._attr_extra_state_attributes: Mapping[str, Any] = {}
         self._attr_hostname: str | None = hostname
         self._attr_ip_address: str | None = None
         self._attr_mac_address: str | None = mac
         self._attr_source_type: SourceType = SourceType.ROUTER
-        self._available: bool = (
-            False  # Move this to OPNsenseEntity once all entity-types are updated
-        )
-
-    # Move this to OPNsenseEntity once all entity-types are updated
-    @property
-    def available(self) -> bool:
-        return self._available
 
     @property
     def source_type(self) -> SourceType:
@@ -179,7 +172,7 @@ def hostname(self) -> str | None:
 
     @property
     def unique_id(self) -> str | None:
-        return self._attr_unique_id
+        return self._attr_mac_address
 
     @property
     def entity_registry_enabled_default(self) -> bool:
@@ -346,6 +339,5 @@ async def _restore_last_state(self) -> None:
                 pass
 
     async def async_added_to_hass(self) -> None:
-        await super().async_added_to_hass()
         await self._restore_last_state()
-        self._handle_coordinator_update()
+        await super().async_added_to_hass()

custom_components/opnsense/helpers.py

@@ -10,7 +10,7 @@ def dict_get(data: Mapping[str, Any], path: str, default=None):
         try:
             key: int | str = int(key) if key.isnumeric() else key
             result = result[key]
-        except:
+        except (TypeError, KeyError, AttributeError):
             result = default
             break
 

custom_components/opnsense/pyopnsense/__init__.py

@@ -35,7 +35,7 @@ def dict_get(data: Mapping[str, Any], path: str, default=None):
         try:
             key = int(key) if key.isnumeric() else key
             result = result[key]
-        except:
+        except (TypeError, KeyError, AttributeError):
             result = default
             break
 
@@ -339,7 +339,7 @@ async def get_system_info(self) -> Mapping[str, Any]:
             if awesomeversion.AwesomeVersion(firmware) < awesomeversion.AwesomeVersion(
                 "24.7"
             ):
-                _LOGGER.info(f"Using legacy get_system_info method for OPNsense < 24.7")
+                _LOGGER.info("Using legacy get_system_info method for OPNsense < 24.7")
                 return await self._get_system_info_legacy()
         except awesomeversion.exceptions.AwesomeVersionCompareException:
             pass
@@ -393,7 +393,8 @@ async def get_firmware_update_info(self):
         # {'status_msg': 'Firmware status check was aborted internally. Please try again.', 'status': 'error'}
         # error could be because data has not been refreshed at all OR an upgrade is currently in progress
         if (
-            status["status"] == "error"
+            not isinstance(status, Mapping)
+            or status.get("status", None) == "error"
             or "last_check" not in status.keys()
             or not isinstance(dict_get(status, "product.product_check"), dict)
             or dict_get(status, "product.product_check") is None
@@ -419,13 +420,13 @@ async def get_firmware_update_info(self):
             if stale:
                 upgradestatus = await self._get("/api/core/firmware/upgradestatus")
                 # print(upgradestatus)
-                if "status" in upgradestatus.keys():
+                if "status" in upgradestatus:
                     # status = running (package refresh in progress OR upgrade in progress)
                     # status = done (refresh/upgrade done)
                     if upgradestatus["status"] == "done":
                         # tigger repo update
                         # should this be /api/core/firmware/upgrade
-                        check = await self._post("/api/core/firmware/check")
+                        # check = await self._post("/api/core/firmware/check")
                         # print(check)
                         refresh_triggered = True
                     else:
@@ -791,7 +792,7 @@ async def get_telemetry(self) -> Mapping[str, Any]:
             if awesomeversion.AwesomeVersion(firmware) < awesomeversion.AwesomeVersion(
                 "24.7"
             ):
-                _LOGGER.info(f"Using legacy get_telemetry method for OPNsense < 24.7")
+                _LOGGER.info("Using legacy get_telemetry method for OPNsense < 24.7")
                 return await self._get_telemetry_legacy()
         except awesomeversion.exceptions.AwesomeVersionCompareException:
             pass
@@ -1383,7 +1384,7 @@ async def get_unbound_blocklist(self) -> Mapping[str, Any]:
         response: Mapping[str, Any] | list = await self._get(
             "/api/unbound/settings/get"
         )
-        if response is None or not isinstance(response, Mapping):
+        if not isinstance(response, Mapping):
             _LOGGER.error("Invalid data returned from get_unbound_blocklist")
             return {}
         # _LOGGER.debug(f"[get_unbound_blocklist] response: {response}")
@@ -1395,7 +1396,7 @@ async def get_unbound_blocklist(self) -> Mapping[str, Any]:
         for attr in ["enabled", "safesearch", "nxdomain", "address"]:
             dnsbl[attr] = dnsbl_settings.get(attr, "")
         for attr in ["type", "lists", "whitelists", "blocklists", "wildcards"]:
-            if isinstance(dnsbl_settings[attr], Mapping):
+            if isinstance(dnsbl_settings.get(attr, None), Mapping):
                 dnsbl[attr] = ",".join(
                     [
                         key

custom_components/opnsense/sensor.py

@@ -19,7 +19,6 @@
     UnitOfInformation,
     UnitOfTemperature,
     UnitOfTime,
-    __version__,
 )
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import entity_platform
@@ -425,20 +424,6 @@ def __init__(
         self._attr_entity_registry_enabled_default: bool = enabled_default
         self._previous_value = None
         self._attr_native_value = None
-        self._attr_extra_state_attributes: Mapping[str, Any] = {}
-        self._available: bool = (
-            False  # Move this to OPNsenseEntity once all entity-types are updated
-        )
-
-    # Move this to OPNsenseEntity once all entity-types are updated
-    @property
-    def available(self) -> bool:
-        return self._available
-
-    # Move this to OPNsenseEntity once all entity-types are updated
-    async def async_added_to_hass(self) -> None:
-        await super().async_added_to_hass()
-        self._handle_coordinator_update()
 
 
 class OPNsenseStaticKeySensor(OPNsenseSensor):
@@ -510,7 +495,7 @@ def _handle_coordinator_update(self) -> None:
         filesystem = self._opnsense_get_filesystem()
         try:
             self._attr_native_value = filesystem["capacity"].strip("%")
-        except (TypeError, KeyError):
+        except (TypeError, KeyError, AttributeError):
             self._available = False
             return
         self._available = True

custom_components/opnsense/switch.py

@@ -262,20 +262,6 @@ def __init__(
         )
         self.entity_description = entity_description
         self._attr_is_on: bool = False
-        self._attr_extra_state_attributes: Mapping[str, Any] = {}
-        self._available: bool = (
-            False  # Move this to OPNsenseEntity once all entity-types are updated
-        )
-
-    # Move this to OPNsenseEntity once all entity-types are updated
-    @property
-    def available(self) -> bool:
-        return self._available
-
-    # Move this to OPNsenseEntity once all entity-types are updated
-    async def async_added_to_hass(self) -> None:
-        await super().async_added_to_hass()
-        self._handle_coordinator_update()
 
 
 class OPNsenseFilterSwitch(OPNsenseSwitch):
@@ -315,7 +301,7 @@ def _handle_coordinator_update(self) -> None:
         self._rule = self._opnsense_get_rule()
         try:
             self._attr_is_on = bool(self._rule.get("disabled", "0") != "1")
-        except (TypeError, KeyError):
+        except (TypeError, KeyError, AttributeError):
             self._available = False
             return
         self._available = True
@@ -390,7 +376,7 @@ def _handle_coordinator_update(self) -> None:
         self._rule = self._opnsense_get_rule()
         try:
             self._attr_is_on = "disabled" not in self._rule
-        except (TypeError, KeyError):
+        except (TypeError, KeyError, AttributeError):
             self._available = False
             return
         self._available = True
@@ -464,7 +450,7 @@ def _handle_coordinator_update(self) -> None:
         self._service = self._opnsense_get_service()
         try:
             self._attr_is_on = self._service[self._prop_name]
-        except (TypeError, KeyError):
+        except (TypeError, KeyError, AttributeError):
             self._available = False
             return
         self._available = True