updated apps

custom_components/powercalc/discovery.py

@@ -51,35 +51,38 @@ _LOGGER = logging.getLogger(__name__)
 _DiscoverySourceT = TypeVar("_DiscoverySourceT", er.RegistryEntry, dr.DeviceEntry)
 
 
-async def get_power_profile_by_source_entity(hass: HomeAssistant, source_entity: SourceEntity) -> PowerProfile | None:
-    """Given a certain entity, lookup the manufacturer and model and return the power profile."""
+def get_discovery_manager(hass: HomeAssistant) -> DiscoveryManager:
+    """Return the shared discovery manager, creating a throwaway one when not yet set up."""
     try:
-        discovery_manager: DiscoveryManager = hass.data[DOMAIN][DATA_DISCOVERY_MANAGER]
+        return hass.data[DOMAIN][DATA_DISCOVERY_MANAGER]  # type: ignore[no-any-return]
     except KeyError:
-        discovery_manager = DiscoveryManager(hass, {})
+        return DiscoveryManager(hass, {})
+
+
+async def _get_power_profile_by_source(
+    hass: HomeAssistant,
+    source_entity: SourceEntity,
+    discovery_by: DiscoveryBy,
+) -> PowerProfile | None:
+    """Look up a power profile for a source entity, discovered either by entity or by device."""
+    discovery_manager = get_discovery_manager(hass)
     model_info = await discovery_manager.extract_model_info_from_device_info(source_entity.entity_entry)
     if not model_info:
         return None
-    profiles = await discovery_manager.find_power_profiles(model_info, source_entity, DiscoveryBy.ENTITY)
+    profiles = await discovery_manager.find_power_profiles(model_info, source_entity, discovery_by)
     return profiles[0] if profiles else None
 
 
+async def get_power_profile_by_source_entity(hass: HomeAssistant, source_entity: SourceEntity) -> PowerProfile | None:
+    """Given a certain entity, lookup the manufacturer and model and return the power profile."""
+    return await _get_power_profile_by_source(hass, source_entity, DiscoveryBy.ENTITY)
+
+
 async def get_power_profile_by_source_device(hass: HomeAssistant, source_entity: SourceEntity) -> PowerProfile | None:
     """Look up a device-discovered power profile for a source entity's device."""
     if not source_entity.device_entry or not source_entity.entity_entry:
         return None
-
-    try:
-        discovery_manager: DiscoveryManager = hass.data[DOMAIN][DATA_DISCOVERY_MANAGER]
-    except KeyError:
-        discovery_manager = DiscoveryManager(hass, {})
-
-    model_info = await discovery_manager.extract_model_info_from_device_info(source_entity.entity_entry)
-    if not model_info:
-        return None
-
-    profiles = await discovery_manager.find_power_profiles(model_info, source_entity, DiscoveryBy.DEVICE)
-    return profiles[0] if profiles else None
+    return await _get_power_profile_by_source(hass, source_entity, DiscoveryBy.DEVICE)
 
 
 class DiscoveryStatus(StrEnum):