From bf1bddba831ee90b9134abb9dfef67e85537bbe9 Mon Sep 17 00:00:00 2001 From: Bram Date: Sun, 5 Jan 2025 10:17:28 +0100 Subject: [PATCH 1/2] feat: optimize discovery and improve logging --- custom_components/powercalc/discovery.py | 43 ++++++++------- .../powercalc/group_include/filter.py | 8 +-- .../powercalc/group_include/include.py | 3 ++ tests/test_discovery.py | 53 +++++++++++++++++++ 4 files changed, 85 insertions(+), 22 deletions(-) diff --git a/custom_components/powercalc/discovery.py b/custom_components/powercalc/discovery.py index 752e95df3..d4ecb2724 100644 --- a/custom_components/powercalc/discovery.py +++ b/custom_components/powercalc/discovery.py @@ -32,11 +32,11 @@ MANUFACTURER_WLED, CalculationStrategy, ) -from .group_include.filter import CategoryFilter, CompositeFilter, FilterOperator, LambdaFilter, NotFilter, get_filtered_entity_list +from .group_include.filter import CategoryFilter, CompositeFilter, DomainFilter, FilterOperator, LambdaFilter, NotFilter, get_filtered_entity_list from .helpers import get_or_create_unique_id from .power_profile.factory import get_power_profile from .power_profile.library import ModelInfo, ProfileLibrary -from .power_profile.power_profile import DiscoveryBy, PowerProfile +from .power_profile.power_profile import DEVICE_TYPE_DOMAIN, DiscoveryBy, PowerProfile _LOGGER = logging.getLogger(__name__) @@ -94,7 +94,10 @@ async def start_discovery(self) -> None: _LOGGER.debug("Start auto discovery") + _LOGGER.debug("Start entity discovery") await self.perform_discovery(self.get_entities, self.create_entity_source, DiscoveryBy.ENTITY) # type: ignore[arg-type] + + _LOGGER.debug("Start device discovery") await self.perform_discovery(self.get_devices, self.create_device_source, DiscoveryBy.DEVICE) # type: ignore[arg-type] _LOGGER.debug("Done auto discovery") @@ -124,6 +127,7 @@ async def perform_discovery( ) -> None: """Generalized discovery procedure for entities and devices.""" for source in await source_provider(): + log_identifier = source.entity_id if discovery_type == DiscoveryBy.ENTITY else source.id try: model_info = await self.extract_model_info_from_device_info(source) if not model_info: @@ -133,10 +137,7 @@ async def perform_discovery( power_profiles = await self.discover_entity(source_entity, model_info, discovery_type) if not power_profiles: - _LOGGER.debug( - "%s: Model not found in library, skipping discovery", - source_entity.entity_id, - ) + _LOGGER.debug("%s: Model not found in library, skipping discovery", log_identifier) continue unique_id = self.create_unique_id( @@ -146,18 +147,16 @@ async def perform_discovery( ) if self._is_already_discovered(source_entity, unique_id): - _LOGGER.debug( - "%s: Already setup with discovery, skipping", - source_entity.entity_id, - ) + _LOGGER.debug("%s: Already setup with discovery, skipping", log_identifier) continue - self._init_entity_discovery(model_info, unique_id, source_entity, power_profiles, {}) - except Exception: # noqa: BLE001 + self._init_entity_discovery(model_info, unique_id, source_entity, log_identifier, power_profiles, {}) + except Exception as err: # noqa: BLE001 _LOGGER.error( - "Error during %s discovery: %s", + "%s: Error during %s discovery: %s", + log_identifier, discovery_type, - source, + err, ) async def discover_entity( @@ -238,6 +237,7 @@ async def init_wled_flow(self, model_info: ModelInfo, source_entity: SourceEntit model_info, unique_id, source_entity, + source_entity.entity_id, power_profiles=None, extra_discovery_data={ CONF_MODE: CalculationStrategy.WLED, @@ -278,6 +278,7 @@ def _check_already_configured(entity: er.RegistryEntry) -> bool: LambdaFilter(lambda entity: entity.device_id is None), LambdaFilter(lambda entity: entity.platform == "mqtt" and "segment" in entity.entity_id), LambdaFilter(lambda entity: entity.platform == "powercalc"), + NotFilter(DomainFilter(DEVICE_TYPE_DOMAIN.values())), ], FilterOperator.OR, ) @@ -287,11 +288,16 @@ async def get_devices(self) -> list: """Fetch device entries.""" return list(dr.async_get(self.hass).devices.values()) - async def extract_model_info_from_device_info(self, entry: er.RegistryEntry | dr.DeviceEntry | None) -> ModelInfo | None: + async def extract_model_info_from_device_info( + self, + entry: er.RegistryEntry | dr.DeviceEntry | None, + ) -> ModelInfo | None: """Try to auto discover manufacturer and model from the known device information.""" if not entry: return None + log_identifier = entry.entity_id if isinstance(entry, er.RegistryEntry) else entry.id + if isinstance(entry, er.RegistryEntry): model_info = await self.get_model_information_from_entity(entry) else: @@ -299,7 +305,7 @@ async def extract_model_info_from_device_info(self, entry: er.RegistryEntry | dr if not model_info: _LOGGER.debug( "%s: Cannot autodiscover model, manufacturer or model unknown from device registry", - entry.id, + log_identifier, ) return None @@ -315,7 +321,7 @@ async def extract_model_info_from_device_info(self, entry: er.RegistryEntry | dr _LOGGER.debug( "%s: Found model information on device (manufacturer=%s, model=%s, model_id=%s)", - entry.id, + log_identifier, model_info.manufacturer, model_info.model, model_info.model_id, @@ -354,6 +360,7 @@ def _init_entity_discovery( model_info: ModelInfo, unique_id: str, source_entity: SourceEntity, + log_identifier: str, power_profiles: list[PowerProfile] | None, extra_discovery_data: dict | None, ) -> None: @@ -384,7 +391,7 @@ def _init_entity_discovery( if source_entity.entity_id != DUMMY_ENTITY_ID: self.initialized_flows.add(source_entity.entity_id) - _LOGGER.debug("%s: Initiating discovery flow, unique_id=%s", source_entity.entity_id, unique_id) + _LOGGER.debug("%s: Initiating discovery flow, unique_id=%s", log_identifier, unique_id) discovery_flow.async_create_flow( self.hass, diff --git a/custom_components/powercalc/group_include/filter.py b/custom_components/powercalc/group_include/filter.py index db8f3ddd1..0e76b2f2f 100644 --- a/custom_components/powercalc/group_include/filter.py +++ b/custom_components/powercalc/group_include/filter.py @@ -1,7 +1,7 @@ from __future__ import annotations import re -from collections.abc import Callable +from collections.abc import Callable, Iterable from enum import StrEnum from typing import Protocol, cast @@ -109,11 +109,11 @@ def is_valid(self, entity: RegistryEntry) -> bool: class DomainFilter(EntityFilter): - def __init__(self, domain: str | list) -> None: - self.domain = domain + def __init__(self, domain: str | Iterable[str]) -> None: + self.domain = domain if isinstance(domain, str) else set(domain) def is_valid(self, entity: RegistryEntry) -> bool: - if isinstance(self.domain, list): + if isinstance(self.domain, set): return entity.domain in self.domain return entity.domain == self.domain diff --git a/custom_components/powercalc/group_include/include.py b/custom_components/powercalc/group_include/include.py index f51bd232c..272ee1aae 100644 --- a/custom_components/powercalc/group_include/include.py +++ b/custom_components/powercalc/group_include/include.py @@ -32,6 +32,9 @@ async def find_entities( resolved_entities: list[Entity] = [] discoverable_entities: list[str] = [] + # new_filter = DomainFilter(DEVICE_TYPE_DOMAIN.values()) + # if entity_filter: + # new_filter = CompositeFilter([new_filter, entity_filter], FilterOperator.OR) source_entities = await get_filtered_entity_list(hass, entity_filter or NullFilter()) if _LOGGER.isEnabledFor(logging.DEBUG): # pragma: no cover _LOGGER.debug("Found possible include entities: %s", [entity.entity_id for entity in source_entities]) diff --git a/tests/test_discovery.py b/tests/test_discovery.py index 9898cb730..8c38d3bb1 100644 --- a/tests/test_discovery.py +++ b/tests/test_discovery.py @@ -775,3 +775,56 @@ async def test_powercalc_sensors_are_ignored_for_discovery( mock_calls = mock_flow_init.mock_calls assert len(mock_calls) == 1 + + +@pytest.mark.parametrize( + "entity_entries,expected_entities", + [ + ( + [ + RegistryEntry( + entity_id="switch.test", + unique_id="1111", + platform="hue", + device_id="hue-device", + ), + ], + ["switch.test"], + ), + # Entity domains that are not supported must be ignored + ( + [ + RegistryEntry( + entity_id="scene.test", + unique_id="1111", + platform="hue", + device_id="hue-device", + ), + RegistryEntry( + entity_id="event.test", + unique_id="2222", + platform="hue", + device_id="hue-device", + ), + ], + [], + ), + # Powercalc sensors should not be considered for discovery + ( + [ + RegistryEntry( + entity_id="sensor.test", + unique_id="1111", + platform="powercalc", + device_id="some-device", + ), + ], + [], + ), + ], +) +async def test_get_entities(hass: HomeAssistant, entity_entries: list[RegistryEntry], expected_entities: list[str]) -> None: + mock_registry(hass, {entity_entry.entity_id: entity_entry for entity_entry in entity_entries}) + discovery_manager = DiscoveryManager(hass, {}) + entity_ids = [entity.entity_id for entity in await discovery_manager.get_entities()] + assert entity_ids == expected_entities From b102a7b8d9ac958bf646584631f3524779621bfd Mon Sep 17 00:00:00 2001 From: Bram Date: Sun, 5 Jan 2025 10:39:18 +0100 Subject: [PATCH 2/2] feat: improve include discovery logic --- .../powercalc/group_include/include.py | 27 +++++++---- tests/group_include/test_include.py | 45 +++++++++++++++++++ 2 files changed, 64 insertions(+), 8 deletions(-) diff --git a/custom_components/powercalc/group_include/include.py b/custom_components/powercalc/group_include/include.py index 272ee1aae..89d3fa5f0 100644 --- a/custom_components/powercalc/group_include/include.py +++ b/custom_components/powercalc/group_include/include.py @@ -12,10 +12,11 @@ DOMAIN, ) from custom_components.powercalc.discovery import get_power_profile_by_source_entity +from custom_components.powercalc.power_profile.power_profile import DEVICE_TYPE_DOMAIN from custom_components.powercalc.sensors.energy import RealEnergySensor from custom_components.powercalc.sensors.power import RealPowerSensor -from .filter import EntityFilter, NullFilter, get_filtered_entity_list +from .filter import CompositeFilter, DomainFilter, EntityFilter, LambdaFilter, get_filtered_entity_list _LOGGER = logging.getLogger(__name__) @@ -32,15 +33,12 @@ async def find_entities( resolved_entities: list[Entity] = [] discoverable_entities: list[str] = [] - # new_filter = DomainFilter(DEVICE_TYPE_DOMAIN.values()) - # if entity_filter: - # new_filter = CompositeFilter([new_filter, entity_filter], FilterOperator.OR) - source_entities = await get_filtered_entity_list(hass, entity_filter or NullFilter()) + source_entities = await get_filtered_entity_list(hass, _build_filter(entity_filter)) if _LOGGER.isEnabledFor(logging.DEBUG): # pragma: no cover _LOGGER.debug("Found possible include entities: %s", [entity.entity_id for entity in source_entities]) - source_entity_powercalc_entity_map: dict[str, list] = domain_data[DATA_CONFIGURED_ENTITIES] - powercalc_entities: dict[str, Entity] = domain_data[DATA_ENTITIES] + source_entity_powercalc_entity_map: dict[str, list] = domain_data.get(DATA_CONFIGURED_ENTITIES, {}) + powercalc_entities: dict[str, Entity] = domain_data.get(DATA_ENTITIES, {}) for source_entity in source_entities: if source_entity.entity_id in source_entity_powercalc_entity_map: resolved_entities.extend(source_entity_powercalc_entity_map[source_entity.entity_id]) @@ -54,7 +52,7 @@ async def find_entities( device_class = source_entity.device_class or source_entity.original_device_class if device_class == SensorDeviceClass.POWER: resolved_entities.append(RealPowerSensor(source_entity.entity_id, source_entity.unit_of_measurement)) - elif device_class == SensorDeviceClass.ENERGY and source_entity.platform != "utility_meter": + elif device_class == SensorDeviceClass.ENERGY: resolved_entities.append(RealEnergySensor(source_entity.entity_id)) power_profile = await get_power_profile_by_source_entity( @@ -65,3 +63,16 @@ async def find_entities( discoverable_entities.append(source_entity.entity_id) return resolved_entities, discoverable_entities + + +def _build_filter(entity_filter: EntityFilter | None) -> EntityFilter: + base_filter = CompositeFilter( + [ + DomainFilter(DEVICE_TYPE_DOMAIN.values()), + LambdaFilter(lambda entity: entity.platform != "utility_meter"), + ], + ) + if not entity_filter: + return base_filter + + return CompositeFilter([base_filter, entity_filter]) diff --git a/tests/group_include/test_include.py b/tests/group_include/test_include.py index be270e3d8..3e2eba55c 100644 --- a/tests/group_include/test_include.py +++ b/tests/group_include/test_include.py @@ -49,6 +49,7 @@ ENTRY_DATA_POWER_ENTITY, SensorType, ) +from custom_components.powercalc.group_include.include import find_entities from custom_components.test.light import MockLight from tests.common import ( create_discoverable_light, @@ -1287,6 +1288,50 @@ async def test_include_logs_warning(hass: HomeAssistant, caplog: pytest.LogCaptu assert "Could not resolve any entities in group" in caplog.text +async def test_irrelevant_entity_domains_are_skipped(hass: HomeAssistant, caplog: pytest.LogCaptureFixture) -> None: + caplog.set_level(logging.DEBUG) + + mock_device_registry( + hass, + { + "device-a": DeviceEntry( + id="device-a", + manufacturer="Signify", + model="LCT012", + ), + }, + ) + mock_registry( + hass, + { + "light.test": RegistryEntry( + entity_id="light.test", + unique_id="2222", + platform="hue", + device_id="device-a", + ), + "scene.test": RegistryEntry( + entity_id="scene.test", + unique_id="3333", + platform="hue", + device_id="device-a", + ), + "event.test": RegistryEntry( + entity_id="event.test", + unique_id="4444", + platform="hue", + device_id="device-a", + ), + }, + ) + _, discoverable_entities = await find_entities(hass) + assert len(discoverable_entities) == 1 + assert "light.test" in discoverable_entities + + assert "scene.test" not in caplog.text + assert "event.test" not in caplog.text + + def _create_powercalc_config_entry( hass: HomeAssistant, source_entity_id: str,