From 9da155955ab1b1b3d20a64f495ae3d18e17b9bb8 Mon Sep 17 00:00:00 2001
From: Tim <tim@swipejobs.com>
Date: Thu, 20 Jul 2023 16:35:26 +1000
Subject: [PATCH] Transport NSW: Set DeviceClass and StateClass (#96928)

* 2023.7.16
- Fix bug with values defaulting to "n/a" in stead of None

* 2023.7.16
- Set device class and state classes on entities

* 2023.7.16
- Set StateClass and DeviceClass directly on the entitiy

* 2023.7.16
- Fix black and ruff issues

* 2023.7.17
- Update logic catering for the 'n/a' response on an API failure
- Add testcase

* - Fix bug in formatting

* 2023.7.17
- Refacotr to consider the "n/a" response returned from the Python lib
on an error or faliure
- Remove setting of StateClass and DeviceClass as requested
- Add "n/a" test case

* 2023.7.17
- Remove unused imports

* 2023.7.18
- Apply review requested changes

* - Additional review change resolved

* Add State and Device class attributes
---
 homeassistant/components/transport_nsw/sensor.py | 4 ++++
 tests/components/transport_nsw/test_sensor.py    | 6 ++++++
 2 files changed, 10 insertions(+)

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):