Netgear add traffic sensors (#65645)

* add sensors

* use entity discription

* use lambda functions

* use seperate coordinators

* styling

* revieuw comments

* set proper default

* move api lock

* fix styling

* Update homeassistant/components/netgear/sensor.py

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

* Update homeassistant/components/netgear/sensor.py

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

* use coordinator data

* fix styling

* fix typing

* move typing

* fix lock

* Update homeassistant/components/netgear/sensor.py

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

* Update homeassistant/components/netgear/sensor.py

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

homeassistant/components/netgear/__init__.py

@@ -9,7 +9,13 @@ from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from .const import DOMAIN, KEY_COORDINATOR, KEY_ROUTER, PLATFORMS
+from .const import (
+    DOMAIN,
+    KEY_COORDINATOR,
+    KEY_COORDINATOR_TRAFFIC,
+    KEY_ROUTER,
+    PLATFORMS,
+)
 from .errors import CannotLoginException
 from .router import NetgearRouter
 
@@ -53,28 +59,41 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         name=router.device_name,
         model=router.model,
         sw_version=router.firmware_version,
+        hw_version=router.hardware_version,
         configuration_url=f"http://{entry.data[CONF_HOST]}/",
     )
 
-    async def async_update_data() -> bool:
+    async def async_update_devices() -> bool:
         """Fetch data from the router."""
-        data = await router.async_update_device_trackers()
+        return await router.async_update_device_trackers()
-        return data
 
-    # Create update coordinator
+    async def async_update_traffic_meter() -> dict:
+        """Fetch data from the router."""
+        return await router.async_get_traffic_meter()
+
+    # Create update coordinators
     coordinator = DataUpdateCoordinator(
         hass,
         _LOGGER,
-        name=router.device_name,
+        name=f"{router.device_name} Devices",
-        update_method=async_update_data,
+        update_method=async_update_devices,
+        update_interval=SCAN_INTERVAL,
+    )
+    coordinator_traffic_meter = DataUpdateCoordinator(
+        hass,
+        _LOGGER,
+        name=f"{router.device_name} Traffic meter",
+        update_method=async_update_traffic_meter,
         update_interval=SCAN_INTERVAL,
     )
 
     await coordinator.async_config_entry_first_refresh()
+    await coordinator_traffic_meter.async_config_entry_first_refresh()
 
     hass.data[DOMAIN][entry.entry_id] = {
         KEY_ROUTER: router,
         KEY_COORDINATOR: coordinator,
+        KEY_COORDINATOR_TRAFFIC: coordinator_traffic_meter,
     }
 
     hass.config_entries.async_setup_platforms(entry, PLATFORMS)

homeassistant/components/netgear/const.py

@@ -9,6 +9,7 @@ CONF_CONSIDER_HOME = "consider_home"
 
 KEY_ROUTER = "router"
 KEY_COORDINATOR = "coordinator"
+KEY_COORDINATOR_TRAFFIC = "coordinator_traffic"
 
 PLATFORMS = [Platform.DEVICE_TRACKER, Platform.SENSOR]
 

homeassistant/components/netgear/router.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 from abc import abstractmethod
+import asyncio
 from datetime import timedelta
 import logging
 
@@ -69,9 +70,11 @@ class NetgearRouter:
         self._password = entry.data[CONF_PASSWORD]
 
         self._info = None
-        self.model = None
+        self.model = ""
-        self.device_name = None
+        self.device_name = ""
-        self.firmware_version = None
+        self.firmware_version = ""
+        self.hardware_version = ""
+        self.serial_number = ""
 
         self.method_version = 1
         consider_home_int = entry.options.get(
@@ -80,7 +83,7 @@ class NetgearRouter:
         self._consider_home = timedelta(seconds=consider_home_int)
 
         self._api: Netgear = None
-        self._attrs = {}
+        self._api_lock = asyncio.Lock()
 
         self.devices = {}
 
@@ -101,6 +104,8 @@ class NetgearRouter:
         self.device_name = self._info.get("DeviceName", DEFAULT_NAME)
         self.model = self._info.get("ModelName")
         self.firmware_version = self._info.get("Firmwareversion")
+        self.hardware_version = self._info.get("Hardwareversion")
+        self.serial_number = self._info["SerialNumber"]
 
         for model in MODELS_V2:
             if self.model.startswith(model):
@@ -149,11 +154,15 @@ class NetgearRouter:
     async def async_get_attached_devices(self) -> list:
         """Get the devices connected to the router."""
         if self.method_version == 1:
-            return await self.hass.async_add_executor_job(
+            async with self._api_lock:
-                self._api.get_attached_devices
+                return await self.hass.async_add_executor_job(
-            )
+                    self._api.get_attached_devices
+                )
 
-        return await self.hass.async_add_executor_job(self._api.get_attached_devices_2)
+        async with self._api_lock:
+            return await self.hass.async_add_executor_job(
+                self._api.get_attached_devices_2
+            )
 
     async def async_update_device_trackers(self, now=None) -> None:
         """Update Netgear devices."""
@@ -186,6 +195,11 @@ class NetgearRouter:
 
         return new_device
 
+    async def async_get_traffic_meter(self) -> None:
+        """Get the traffic meter data of the router."""
+        async with self._api_lock:
+            return await self.hass.async_add_executor_job(self._api.get_traffic_meter)
+
     @property
     def port(self) -> int:
         """Port used by the API."""
@@ -261,3 +275,44 @@ class NetgearDeviceEntity(NetgearBaseEntity):
             default_model=self._device["device_model"],
             via_device=(DOMAIN, self._router.unique_id),
         )
+
+
+class NetgearRouterEntity(CoordinatorEntity):
+    """Base class for a Netgear router entity."""
+
+    def __init__(
+        self, coordinator: DataUpdateCoordinator, router: NetgearRouter
+    ) -> None:
+        """Initialize a Netgear device."""
+        super().__init__(coordinator)
+        self._router = router
+        self._name = router.device_name
+        self._unique_id = router.serial_number
+
+    @abstractmethod
+    @callback
+    def async_update_device(self) -> None:
+        """Update the Netgear device."""
+
+    @callback
+    def _handle_coordinator_update(self) -> None:
+        """Handle updated data from the coordinator."""
+        self.async_update_device()
+        super()._handle_coordinator_update()
+
+    @property
+    def unique_id(self) -> str:
+        """Return a unique ID."""
+        return self._unique_id
+
+    @property
+    def name(self) -> str:
+        """Return the name."""
+        return self._name
+
+    @property
+    def device_info(self) -> DeviceInfo:
+        """Return the device information."""
+        return DeviceInfo(
+            identifiers={(DOMAIN, self._router.unique_id)},
+        )

homeassistant/components/netgear/sensor.py

@@ -1,30 +1,35 @@
 """Support for Netgear routers."""
+from collections.abc import Callable
+from dataclasses import dataclass
+
 from homeassistant.components.sensor import (
     SensorDeviceClass,
     SensorEntity,
     SensorEntityDescription,
 )
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import PERCENTAGE
+from homeassistant.const import DATA_MEGABYTES, PERCENTAGE
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from .const import DOMAIN, KEY_COORDINATOR, KEY_ROUTER
+from .const import DOMAIN, KEY_COORDINATOR, KEY_COORDINATOR_TRAFFIC, KEY_ROUTER
-from .router import NetgearDeviceEntity, NetgearRouter
+from .router import NetgearDeviceEntity, NetgearRouter, NetgearRouterEntity
 
 SENSOR_TYPES = {
     "type": SensorEntityDescription(
         key="type",
         name="link type",
         entity_category=EntityCategory.DIAGNOSTIC,
+        icon="mdi:lan",
     ),
     "link_rate": SensorEntityDescription(
         key="link_rate",
         name="link rate",
         native_unit_of_measurement="Mbps",
         entity_category=EntityCategory.DIAGNOSTIC,
+        icon="mdi:speedometer",
     ),
     "signal": SensorEntityDescription(
         key="signal",
@@ -37,23 +42,185 @@ SENSOR_TYPES = {
         key="ssid",
         name="ssid",
         entity_category=EntityCategory.DIAGNOSTIC,
+        icon="mdi:wifi-marker",
     ),
     "conn_ap_mac": SensorEntityDescription(
         key="conn_ap_mac",
         name="access point mac",
         entity_category=EntityCategory.DIAGNOSTIC,
+        icon="mdi:router-network",
     ),
 }
 
 
+@dataclass
+class NetgearSensorEntityDescription(SensorEntityDescription):
+    """Class describing Netgear sensor entities."""
+
+    value: Callable = lambda data: data
+    index: int = 0
+
+
+SENSOR_TRAFFIC_TYPES = [
+    NetgearSensorEntityDescription(
+        key="NewTodayUpload",
+        name="Upload today",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        native_unit_of_measurement=DATA_MEGABYTES,
+        icon="mdi:upload",
+    ),
+    NetgearSensorEntityDescription(
+        key="NewTodayDownload",
+        name="Download today",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        native_unit_of_measurement=DATA_MEGABYTES,
+        icon="mdi:download",
+    ),
+    NetgearSensorEntityDescription(
+        key="NewYesterdayUpload",
+        name="Upload yesterday",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        native_unit_of_measurement=DATA_MEGABYTES,
+        icon="mdi:upload",
+    ),
+    NetgearSensorEntityDescription(
+        key="NewYesterdayDownload",
+        name="Download yesterday",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        native_unit_of_measurement=DATA_MEGABYTES,
+        icon="mdi:download",
+    ),
+    NetgearSensorEntityDescription(
+        key="NewWeekUpload",
+        name="Upload week",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        native_unit_of_measurement=DATA_MEGABYTES,
+        icon="mdi:upload",
+        index=0,
+        value=lambda data: data[0] if data is not None else None,
+    ),
+    NetgearSensorEntityDescription(
+        key="NewWeekUpload",
+        name="Upload week average",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        native_unit_of_measurement=DATA_MEGABYTES,
+        icon="mdi:upload",
+        index=1,
+        value=lambda data: data[1] if data is not None else None,
+    ),
+    NetgearSensorEntityDescription(
+        key="NewWeekDownload",
+        name="Download week",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        native_unit_of_measurement=DATA_MEGABYTES,
+        icon="mdi:download",
+        index=0,
+        value=lambda data: data[0] if data is not None else None,
+    ),
+    NetgearSensorEntityDescription(
+        key="NewWeekDownload",
+        name="Download week average",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        native_unit_of_measurement=DATA_MEGABYTES,
+        icon="mdi:download",
+        index=1,
+        value=lambda data: data[1] if data is not None else None,
+    ),
+    NetgearSensorEntityDescription(
+        key="NewMonthUpload",
+        name="Upload month",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        native_unit_of_measurement=DATA_MEGABYTES,
+        icon="mdi:upload",
+        index=0,
+        value=lambda data: data[0] if data is not None else None,
+    ),
+    NetgearSensorEntityDescription(
+        key="NewMonthUpload",
+        name="Upload month average",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        native_unit_of_measurement=DATA_MEGABYTES,
+        icon="mdi:upload",
+        index=1,
+        value=lambda data: data[1] if data is not None else None,
+    ),
+    NetgearSensorEntityDescription(
+        key="NewMonthDownload",
+        name="Download month",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        native_unit_of_measurement=DATA_MEGABYTES,
+        icon="mdi:download",
+        index=0,
+        value=lambda data: data[0] if data is not None else None,
+    ),
+    NetgearSensorEntityDescription(
+        key="NewMonthDownload",
+        name="Download month average",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        native_unit_of_measurement=DATA_MEGABYTES,
+        icon="mdi:download",
+        index=1,
+        value=lambda data: data[1] if data is not None else None,
+    ),
+    NetgearSensorEntityDescription(
+        key="NewLastMonthUpload",
+        name="Upload last month",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        native_unit_of_measurement=DATA_MEGABYTES,
+        icon="mdi:upload",
+        index=0,
+        value=lambda data: data[0] if data is not None else None,
+    ),
+    NetgearSensorEntityDescription(
+        key="NewLastMonthUpload",
+        name="Upload last month average",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        native_unit_of_measurement=DATA_MEGABYTES,
+        icon="mdi:upload",
+        index=1,
+        value=lambda data: data[1] if data is not None else None,
+    ),
+    NetgearSensorEntityDescription(
+        key="NewLastMonthDownload",
+        name="Download last month",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        native_unit_of_measurement=DATA_MEGABYTES,
+        icon="mdi:download",
+        index=0,
+        value=lambda data: data[0] if data is not None else None,
+    ),
+    NetgearSensorEntityDescription(
+        key="NewLastMonthDownload",
+        name="Download last month average",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        native_unit_of_measurement=DATA_MEGABYTES,
+        icon="mdi:download",
+        index=1,
+        value=lambda data: data[1] if data is not None else None,
+    ),
+]
+
+
 async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up device tracker for Netgear component."""
     router = hass.data[DOMAIN][entry.entry_id][KEY_ROUTER]
     coordinator = hass.data[DOMAIN][entry.entry_id][KEY_COORDINATOR]
-    tracked = set()
+    coordinator_traffic = hass.data[DOMAIN][entry.entry_id][KEY_COORDINATOR_TRAFFIC]
 
+    # Router entities
+    router_entities = []
+
+    for description in SENSOR_TRAFFIC_TYPES:
+        router_entities.append(
+            NetgearRouterSensorEntity(coordinator_traffic, router, description)
+        )
+
+    async_add_entities(router_entities)
+
+    # Entities per network device
+    tracked = set()
     sensors = ["type", "link_rate", "signal"]
     if router.method_version == 2:
         sensors.extend(["ssid", "conn_ap_mac"])
@@ -119,3 +286,37 @@ class NetgearSensorEntity(NetgearDeviceEntity, SensorEntity):
         self._active = self._device["active"]
         if self._device.get(self._attribute) is not None:
             self._state = self._device[self._attribute]
+
+
+class NetgearRouterSensorEntity(NetgearRouterEntity, SensorEntity):
+    """Representation of a device connected to a Netgear router."""
+
+    _attr_entity_registry_enabled_default = False
+    entity_description: NetgearSensorEntityDescription
+
+    def __init__(
+        self,
+        coordinator: DataUpdateCoordinator,
+        router: NetgearRouter,
+        entity_description: NetgearSensorEntityDescription,
+    ) -> None:
+        """Initialize a Netgear device."""
+        super().__init__(coordinator, router)
+        self.entity_description = entity_description
+        self._name = f"{router.device_name} {entity_description.name}"
+        self._unique_id = f"{router.serial_number}-{entity_description.key}-{entity_description.index}"
+
+        self._value = None
+        self.async_update_device()
+
+    @property
+    def native_value(self):
+        """Return the state of the sensor."""
+        return self._value
+
+    @callback
+    def async_update_device(self) -> None:
+        """Update the Netgear device."""
+        if self.coordinator.data is not None:
+            data = self.coordinator.data.get(self.entity_description.key)
+            self._value = self.entity_description.value(data)