From be882c5de3735168e913f65635d8fc3358397614 Mon Sep 17 00:00:00 2001 From: Snuffy2 Date: Sun, 29 Sep 2024 20:09:07 -0400 Subject: [PATCH] Optimize entities final cleanup --- custom_components/opnsense/__init__.py | 79 ++++--------------- custom_components/opnsense/binary_sensor.py | 16 ---- custom_components/opnsense/device_tracker.py | 20 ++--- custom_components/opnsense/helpers.py | 2 +- .../opnsense/pyopnsense/__init__.py | 17 ++-- custom_components/opnsense/sensor.py | 17 +--- custom_components/opnsense/switch.py | 20 +---- custom_components/opnsense/update.py | 23 +----- 8 files changed, 39 insertions(+), 155 deletions(-) diff --git a/custom_components/opnsense/__init__.py b/custom_components/opnsense/__init__.py index bf12bf5..c1c3608 100644 --- a/custom_components/opnsense/__init__.py +++ b/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,61 +201,6 @@ 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""" @@ -265,21 +208,28 @@ 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() diff --git a/custom_components/opnsense/binary_sensor.py b/custom_components/opnsense/binary_sensor.py index be3321b..6d4a253 100644 --- a/custom_components/opnsense/binary_sensor.py +++ b/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): diff --git a/custom_components/opnsense/device_tracker.py b/custom_components/opnsense/device_tracker.py index 02822db..c439d88 100644 --- a/custom_components/opnsense/device_tracker.py +++ b/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,7 +137,7 @@ 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 @@ -143,19 +145,10 @@ def __init__( 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() diff --git a/custom_components/opnsense/helpers.py b/custom_components/opnsense/helpers.py index 6c165b6..35f34dc 100644 --- a/custom_components/opnsense/helpers.py +++ b/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 diff --git a/custom_components/opnsense/pyopnsense/__init__.py b/custom_components/opnsense/pyopnsense/__init__.py index a9cbd12..46902e5 100644 --- a/custom_components/opnsense/pyopnsense/__init__.py +++ b/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 @@ -329,7 +329,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 @@ -383,7 +383,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 @@ -409,13 +410,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: @@ -781,7 +782,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 @@ -1373,7 +1374,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}") @@ -1385,7 +1386,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 diff --git a/custom_components/opnsense/sensor.py b/custom_components/opnsense/sensor.py index fd9d932..13e4679 100644 --- a/custom_components/opnsense/sensor.py +++ b/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 diff --git a/custom_components/opnsense/switch.py b/custom_components/opnsense/switch.py index 7175864..8f901a3 100644 --- a/custom_components/opnsense/switch.py +++ b/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 diff --git a/custom_components/opnsense/update.py b/custom_components/opnsense/update.py index f3eb16f..11b7676 100644 --- a/custom_components/opnsense/update.py +++ b/custom_components/opnsense/update.py @@ -1,7 +1,6 @@ """OPNsense integration.""" import asyncio -from collections.abc import Mapping import logging from typing import Any @@ -79,20 +78,6 @@ def __init__( self._attr_latest_version: str | None = None self._attr_release_url: str | None = None self._release_notes: str | None = 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 OPNsenseFirmwareUpdatesAvailableUpdate(OPNsenseUpdate): @@ -104,7 +89,7 @@ def _handle_coordinator_update(self) -> None: if state["firmware_update_info"]["status"] == "error": self._available = False return - except (TypeError, KeyError): + except (TypeError, KeyError, AttributeError): self._available = False return self._available = True @@ -113,7 +98,7 @@ def _handle_coordinator_update(self) -> None: self._attr_installed_version = dict_get( state, "firmware_update_info.product.product_version" ) - except (TypeError, KeyError): + except (TypeError, KeyError, AttributeError): self._attr_installed_version = None try: @@ -138,7 +123,7 @@ def _handle_coordinator_update(self) -> None: ) self._attr_latest_version = product_latest - except (TypeError, KeyError): + except (TypeError, KeyError, AttributeError): self._attr_latest_version = None self._attr_release_url = ( @@ -218,7 +203,7 @@ def _handle_coordinator_update(self) -> None: - reboot needed: {upgrade_needs_reboot} """ - except (TypeError, KeyError): + except (TypeError, KeyError, AttributeError): self._release_notes = None self._release_notes = summary