diff --git a/homeassistant/components/unifi/sensor.py b/homeassistant/components/unifi/sensor.py index 86c6b0d6352..f07269dfdd2 100644 --- a/homeassistant/components/unifi/sensor.py +++ b/homeassistant/components/unifi/sensor.py @@ -27,10 +27,11 @@ from homeassistant.components.sensor import ( SensorDeviceClass, SensorEntity, SensorEntityDescription, + SensorStateClass, UnitOfTemperature, ) from homeassistant.config_entries import ConfigEntry -from homeassistant.const import EntityCategory, UnitOfInformation, UnitOfPower +from homeassistant.const import EntityCategory, UnitOfDataRate, UnitOfPower from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddEntitiesCallback import homeassistant.util.dt as dt_util @@ -133,8 +134,10 @@ class UnifiSensorEntityDescription( ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = ( UnifiSensorEntityDescription[Clients, Client]( key="Bandwidth sensor RX", - entity_category=EntityCategory.DIAGNOSTIC, - native_unit_of_measurement=UnitOfInformation.MEGABYTES, + device_class=SensorDeviceClass.DATA_RATE, + state_class=SensorStateClass.MEASUREMENT, + native_unit_of_measurement=UnitOfDataRate.MEGABYTES_PER_SECOND, + icon="mdi:upload", has_entity_name=True, allowed_fn=lambda controller, _: controller.option_allow_bandwidth_sensors, api_handler_fn=lambda api: api.clients, @@ -151,8 +154,10 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = ( ), UnifiSensorEntityDescription[Clients, Client]( key="Bandwidth sensor TX", - entity_category=EntityCategory.DIAGNOSTIC, - native_unit_of_measurement=UnitOfInformation.MEGABYTES, + device_class=SensorDeviceClass.DATA_RATE, + state_class=SensorStateClass.MEASUREMENT, + native_unit_of_measurement=UnitOfDataRate.MEGABYTES_PER_SECOND, + icon="mdi:download", has_entity_name=True, allowed_fn=lambda controller, _: controller.option_allow_bandwidth_sensors, api_handler_fn=lambda api: api.clients, diff --git a/tests/components/unifi/test_sensor.py b/tests/components/unifi/test_sensor.py index b652c38abdb..f4366b98fc3 100644 --- a/tests/components/unifi/test_sensor.py +++ b/tests/components/unifi/test_sensor.py @@ -8,9 +8,11 @@ import pytest from homeassistant.components.device_tracker import DOMAIN as TRACKER_DOMAIN from homeassistant.components.sensor import ( + ATTR_STATE_CLASS, DOMAIN as SENSOR_DOMAIN, SCAN_INTERVAL, SensorDeviceClass, + SensorStateClass, ) from homeassistant.components.unifi.const import ( CONF_ALLOW_BANDWIDTH_SENSORS, @@ -354,16 +356,28 @@ async def test_bandwidth_sensors( assert len(hass.states.async_all()) == 5 assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 4 - assert hass.states.get("sensor.wired_client_rx").state == "1234.0" - assert hass.states.get("sensor.wired_client_tx").state == "5678.0" - assert hass.states.get("sensor.wireless_client_rx").state == "2345.0" - assert hass.states.get("sensor.wireless_client_tx").state == "6789.0" - ent_reg = er.async_get(hass) - assert ( - ent_reg.async_get("sensor.wired_client_rx").entity_category - is EntityCategory.DIAGNOSTIC - ) + # Verify sensor attributes and state + + wrx_sensor = hass.states.get("sensor.wired_client_rx") + assert wrx_sensor.attributes.get(ATTR_DEVICE_CLASS) == SensorDeviceClass.DATA_RATE + assert wrx_sensor.attributes.get(ATTR_STATE_CLASS) == SensorStateClass.MEASUREMENT + assert wrx_sensor.state == "1234.0" + + wtx_sensor = hass.states.get("sensor.wired_client_tx") + assert wtx_sensor.attributes.get(ATTR_DEVICE_CLASS) == SensorDeviceClass.DATA_RATE + assert wtx_sensor.attributes.get(ATTR_STATE_CLASS) == SensorStateClass.MEASUREMENT + assert wtx_sensor.state == "5678.0" + + wlrx_sensor = hass.states.get("sensor.wireless_client_rx") + assert wlrx_sensor.attributes.get(ATTR_DEVICE_CLASS) == SensorDeviceClass.DATA_RATE + assert wlrx_sensor.attributes.get(ATTR_STATE_CLASS) == SensorStateClass.MEASUREMENT + assert wlrx_sensor.state == "2345.0" + + wltx_sensor = hass.states.get("sensor.wireless_client_tx") + assert wltx_sensor.attributes.get(ATTR_DEVICE_CLASS) == SensorDeviceClass.DATA_RATE + assert wltx_sensor.attributes.get(ATTR_STATE_CLASS) == SensorStateClass.MEASUREMENT + assert wltx_sensor.state == "6789.0" # Verify state update