From 9a8fe0490749e440577c0211b2e460c72baf0904 Mon Sep 17 00:00:00 2001
From: Tim <tim@craibo.com>
Date: Tue, 18 Jul 2023 23:12:43 +1000
Subject: [PATCH] Resolve bugs with Transport NSW (#96692)

* 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
---
 .../components/transport_nsw/sensor.py        | 30 ++++++++++-------
 tests/components/transport_nsw/test_sensor.py | 33 +++++++++++++++++++
 2 files changed, 52 insertions(+), 11 deletions(-)

diff --git a/homeassistant/components/transport_nsw/sensor.py b/homeassistant/components/transport_nsw/sensor.py
index 83fa590429f..0a740ec4347 100644
--- a/homeassistant/components/transport_nsw/sensor.py
+++ b/homeassistant/components/transport_nsw/sensor.py
@@ -6,7 +6,10 @@ from datetime import timedelta
 from TransportNSW import TransportNSW
 import voluptuous as vol
 
-from homeassistant.components.sensor import PLATFORM_SCHEMA, SensorEntity
+from homeassistant.components.sensor import (
+    PLATFORM_SCHEMA,
+    SensorEntity,
+)
 from homeassistant.const import ATTR_MODE, CONF_API_KEY, CONF_NAME, UnitOfTime
 from homeassistant.core import HomeAssistant
 import homeassistant.helpers.config_validation as cv
@@ -121,6 +124,11 @@ class TransportNSWSensor(SensorEntity):
         self._icon = ICONS[self._times[ATTR_MODE]]
 
 
+def _get_value(value):
+    """Replace the API response 'n/a' value with None."""
+    return None if (value is None or value == "n/a") else value
+
+
 class PublicTransportData:
     """The Class for handling the data retrieval."""
 
@@ -132,10 +140,10 @@ class PublicTransportData:
         self._api_key = api_key
         self.info = {
             ATTR_ROUTE: self._route,
-            ATTR_DUE_IN: "n/a",
-            ATTR_DELAY: "n/a",
-            ATTR_REAL_TIME: "n/a",
-            ATTR_DESTINATION: "n/a",
+            ATTR_DUE_IN: None,
+            ATTR_DELAY: None,
+            ATTR_REAL_TIME: None,
+            ATTR_DESTINATION: None,
             ATTR_MODE: None,
         }
         self.tnsw = TransportNSW()
@@ -146,10 +154,10 @@ class PublicTransportData:
             self._stop_id, self._route, self._destination, self._api_key
         )
         self.info = {
-            ATTR_ROUTE: _data["route"],
-            ATTR_DUE_IN: _data["due"],
-            ATTR_DELAY: _data["delay"],
-            ATTR_REAL_TIME: _data["real_time"],
-            ATTR_DESTINATION: _data["destination"],
-            ATTR_MODE: _data["mode"],
+            ATTR_ROUTE: _get_value(_data["route"]),
+            ATTR_DUE_IN: _get_value(_data["due"]),
+            ATTR_DELAY: _get_value(_data["delay"]),
+            ATTR_REAL_TIME: _get_value(_data["real_time"]),
+            ATTR_DESTINATION: _get_value(_data["destination"]),
+            ATTR_MODE: _get_value(_data["mode"]),
         }
diff --git a/tests/components/transport_nsw/test_sensor.py b/tests/components/transport_nsw/test_sensor.py
index 181c5fdd1e4..f9ead2a3054 100644
--- a/tests/components/transport_nsw/test_sensor.py
+++ b/tests/components/transport_nsw/test_sensor.py
@@ -42,3 +42,36 @@ 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"
+
+
+def get_departuresMock_notFound(_stop_id, route, destination, api_key):
+    """Mock TransportNSW departures loading."""
+    data = {
+        "stop_id": "n/a",
+        "route": "n/a",
+        "due": "n/a",
+        "delay": "n/a",
+        "real_time": "n/a",
+        "destination": "n/a",
+        "mode": "n/a",
+    }
+    return data
+
+
+@patch(
+    "TransportNSW.TransportNSW.get_departures", side_effect=get_departuresMock_notFound
+)
+async def test_transportnsw_config_not_found(
+    mocked_get_departures_not_found, hass: HomeAssistant
+) -> None:
+    """Test minimal TransportNSW configuration."""
+    assert await async_setup_component(hass, "sensor", VALID_CONFIG)
+    await hass.async_block_till_done()
+    state = hass.states.get("sensor.next_bus")
+    assert state.state == "unknown"
+    assert state.attributes["stop_id"] == "209516"
+    assert state.attributes["route"] is None
+    assert state.attributes["delay"] is None
+    assert state.attributes["real_time"] is None
+    assert state.attributes["destination"] is None
+    assert state.attributes["mode"] is None