diff --git a/homeassistant/components/transport_nsw/sensor.py b/homeassistant/components/transport_nsw/sensor.py index 0a740ec4347..520b1a5626b 100644 --- a/homeassistant/components/transport_nsw/sensor.py +++ b/homeassistant/components/transport_nsw/sensor.py @@ -8,7 +8,9 @@ import voluptuous as vol from homeassistant.components.sensor import ( PLATFORM_SCHEMA, + SensorDeviceClass, SensorEntity, + SensorStateClass, ) from homeassistant.const import ATTR_MODE, CONF_API_KEY, CONF_NAME, UnitOfTime from homeassistant.core import HomeAssistant @@ -73,6 +75,8 @@ class TransportNSWSensor(SensorEntity): """Implementation of an Transport NSW sensor.""" _attr_attribution = "Data provided by Transport NSW" + _attr_device_class = SensorDeviceClass.DURATION + _attr_state_class = SensorStateClass.MEASUREMENT def __init__(self, data, stop_id, name): """Initialize the sensor.""" diff --git a/tests/components/transport_nsw/test_sensor.py b/tests/components/transport_nsw/test_sensor.py index f9ead2a3054..46aee182b53 100644 --- a/tests/components/transport_nsw/test_sensor.py +++ b/tests/components/transport_nsw/test_sensor.py @@ -1,6 +1,10 @@ """The tests for the Transport NSW (AU) sensor platform.""" from unittest.mock import patch +from homeassistant.components.sensor import ( + SensorDeviceClass, + SensorStateClass, +) from homeassistant.core import HomeAssistant from homeassistant.setup import async_setup_component @@ -42,6 +46,8 @@ async def test_transportnsw_config(mocked_get_departures, hass: HomeAssistant) - assert state.attributes["real_time"] == "y" assert state.attributes["destination"] == "Palm Beach" assert state.attributes["mode"] == "Bus" + assert state.attributes["device_class"] == SensorDeviceClass.DURATION + assert state.attributes["state_class"] == SensorStateClass.MEASUREMENT def get_departuresMock_notFound(_stop_id, route, destination, api_key):