diff --git a/homeassistant/components/fritz/sensor.py b/homeassistant/components/fritz/sensor.py index 5e4b18eebca..7874fb87b00 100644 --- a/homeassistant/components/fritz/sensor.py +++ b/homeassistant/components/fritz/sensor.py @@ -5,7 +5,7 @@ from collections.abc import Callable from dataclasses import dataclass from datetime import datetime, timedelta import logging -from typing import Any, Literal +from typing import Any from fritzconnection.core.exceptions import FritzConnectionException from fritzconnection.lib.fritzstatus import FritzStatus @@ -134,6 +134,15 @@ def _retrieve_link_attenuation_received_state( return status.attenuation[1] / 10 # type: ignore[no-any-return] +@dataclass +class ConnectionInfo: + """Fritz sensor connection information class.""" + + connection: str + mesh_role: MeshRoles + wan_enabled: bool + + @dataclass class FritzRequireKeysMixin: """Fritz sensor data class.""" @@ -145,8 +154,7 @@ class FritzRequireKeysMixin: class FritzSensorEntityDescription(SensorEntityDescription, FritzRequireKeysMixin): """Describes Fritz sensor entity.""" - connection_type: Literal["dsl"] | None = None - exclude_mesh_role: MeshRoles = MeshRoles.SLAVE + is_suitable: Callable[[ConnectionInfo], bool] = lambda info: info.wan_enabled SENSOR_TYPES: tuple[FritzSensorEntityDescription, ...] = ( @@ -162,7 +170,7 @@ SENSOR_TYPES: tuple[FritzSensorEntityDescription, ...] = ( device_class=SensorDeviceClass.TIMESTAMP, entity_category=EntityCategory.DIAGNOSTIC, value_fn=_retrieve_device_uptime_state, - exclude_mesh_role=MeshRoles.NONE, + is_suitable=lambda info: info.mesh_role != MeshRoles.NONE, ), FritzSensorEntityDescription( key="connection_uptime", @@ -225,7 +233,6 @@ SENSOR_TYPES: tuple[FritzSensorEntityDescription, ...] = ( native_unit_of_measurement=DATA_RATE_KILOBITS_PER_SECOND, icon="mdi:upload", value_fn=_retrieve_link_kb_s_sent_state, - connection_type=DSL_CONNECTION, ), FritzSensorEntityDescription( key="link_kb_s_received", @@ -233,7 +240,6 @@ SENSOR_TYPES: tuple[FritzSensorEntityDescription, ...] = ( native_unit_of_measurement=DATA_RATE_KILOBITS_PER_SECOND, icon="mdi:download", value_fn=_retrieve_link_kb_s_received_state, - connection_type=DSL_CONNECTION, ), FritzSensorEntityDescription( key="link_noise_margin_sent", @@ -241,7 +247,7 @@ SENSOR_TYPES: tuple[FritzSensorEntityDescription, ...] = ( native_unit_of_measurement=SIGNAL_STRENGTH_DECIBELS, icon="mdi:upload", value_fn=_retrieve_link_noise_margin_sent_state, - connection_type=DSL_CONNECTION, + is_suitable=lambda info: info.wan_enabled and info.connection == DSL_CONNECTION, ), FritzSensorEntityDescription( key="link_noise_margin_received", @@ -249,7 +255,7 @@ SENSOR_TYPES: tuple[FritzSensorEntityDescription, ...] = ( native_unit_of_measurement=SIGNAL_STRENGTH_DECIBELS, icon="mdi:download", value_fn=_retrieve_link_noise_margin_received_state, - connection_type=DSL_CONNECTION, + is_suitable=lambda info: info.wan_enabled and info.connection == DSL_CONNECTION, ), FritzSensorEntityDescription( key="link_attenuation_sent", @@ -257,7 +263,7 @@ SENSOR_TYPES: tuple[FritzSensorEntityDescription, ...] = ( native_unit_of_measurement=SIGNAL_STRENGTH_DECIBELS, icon="mdi:upload", value_fn=_retrieve_link_attenuation_sent_state, - connection_type=DSL_CONNECTION, + is_suitable=lambda info: info.wan_enabled and info.connection == DSL_CONNECTION, ), FritzSensorEntityDescription( key="link_attenuation_received", @@ -265,7 +271,7 @@ SENSOR_TYPES: tuple[FritzSensorEntityDescription, ...] = ( native_unit_of_measurement=SIGNAL_STRENGTH_DECIBELS, icon="mdi:download", value_fn=_retrieve_link_attenuation_received_state, - connection_type=DSL_CONNECTION, + is_suitable=lambda info: info.wan_enabled and info.connection == DSL_CONNECTION, ), ) @@ -278,19 +284,22 @@ async def async_setup_entry( avm_wrapper: AvmWrapper = hass.data[DOMAIN][entry.entry_id] link_properties = await avm_wrapper.async_get_wan_link_properties() - dsl: bool = link_properties.get("NewWANAccessType") == "DSL" + connection_info = ConnectionInfo( + connection=link_properties.get("NewWANAccessType", "").lower(), + mesh_role=avm_wrapper.mesh_role, + wan_enabled=avm_wrapper.device_is_router, + ) _LOGGER.debug( - "WANAccessType of FritzBox %s is '%s'", + "ConnectionInfo for FritzBox %s: %s", avm_wrapper.host, - link_properties.get("NewWANAccessType"), + connection_info, ) entities = [ FritzBoxSensor(avm_wrapper, entry.title, description) for description in SENSOR_TYPES - if (dsl or description.connection_type != DSL_CONNECTION) - and description.exclude_mesh_role != avm_wrapper.mesh_role + if description.is_suitable(connection_info) ] async_add_entities(entities, True)