only create ipv4_conn_type sensor if supported

AlexandrErohin · Jan 13, 2025 · a9a559a · a9a559a
1 parent 3d9fd19
commit a9a559a
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 7 deletions.
diff --git a/custom_components/tplink_router/coordinator.py b/custom_components/tplink_router/coordinator.py
@@ -81,6 +81,6 @@ async def _async_update_data(self):
         # Only fetch if router is ipv4_status compatible
         if hasattr(self.router, "get_ipv4_status"):
             self.ipv4_status = await self.hass.async_add_executor_job(TPLinkRouterCoordinator.request, self.router,
-                                                                    self.router.get_ipv4_status)
+                                                                      self.router.get_ipv4_status)
         else:
             self.ipv4_status = None
diff --git a/custom_components/tplink_router/sensor.py b/custom_components/tplink_router/sensor.py
@@ -96,13 +96,17 @@ class TPLinkRouterSensorEntityDescription(SensorEntityDescription, TPLinkRouterS
 
 
 async def async_setup_entry(
-        hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     coordinator = hass.data[DOMAIN][entry.entry_id]
 
     sensors = []
 
     for description in SENSOR_TYPES:
+        if description.key == "ipv4_conn_type" and not hasattr(
+            coordinator.router, "get_ipv4_status"
+        ):
+            continue
         sensors.append(TPLinkRouterSensor(coordinator, description))
     async_add_entities(sensors, False)
 
@@ -127,7 +131,7 @@ def __init__(
     @callback
     def _handle_coordinator_update(self) -> None:
         """Handle updated data from the coordinator."""
-        if self.entity_description.key.startswith("ipv4"):
+        if self.entity_description.key == "ipv4_conn_type":
             data = self.coordinator.ipv4_status
         else:
             data = self.coordinator.status
@@ -138,20 +142,18 @@ def _handle_coordinator_update(self) -> None:
 
         # Notify Home Assistant about the state update
         self.async_write_ha_state()
-    
+
     @property
     def available(self) -> bool:
         """Return True if entity is available."""
         if self.entity_description.key.startswith("ipv4"):
-            # Check if the router supports `get_ipv4_status` and if ipv4_status data is available
             return (
                 hasattr(self.coordinator.router, "get_ipv4_status") and
                 self.coordinator.ipv4_status is not None and
                 self.entity_description.value(self.coordinator.ipv4_status) is not None
             )
         else:
-            # General availability check for other sensors
             return (
                 self.coordinator.status is not None and
                 self.entity_description.value(self.coordinator.status) is not None
-        )
+            )