From f59c8d985d85b9d45773bd095be6ab8aa7d344c8 Mon Sep 17 00:00:00 2001
From: Kevin Stillhammer <kevin.stillhammer@gmail.com>
Date: Sat, 17 Sep 2022 06:03:18 +0200
Subject: [PATCH] Correct unit for here_travel_time distance sensor (#78303)

Signed-off-by: Kevin Stillhammer <kevin.stillhammer@gmail.com>

Signed-off-by: Kevin Stillhammer <kevin.stillhammer@gmail.com>
---
 .../components/here_travel_time/sensor.py     | 36 +++++++++++++++----
 .../here_travel_time/test_sensor.py           | 15 +++++++-
 2 files changed, 44 insertions(+), 7 deletions(-)

diff --git a/homeassistant/components/here_travel_time/sensor.py b/homeassistant/components/here_travel_time/sensor.py
index e7bc88f2210..6bb4052dd9f 100644
--- a/homeassistant/components/here_travel_time/sensor.py
+++ b/homeassistant/components/here_travel_time/sensor.py
@@ -23,6 +23,9 @@ from homeassistant.const import (
     CONF_MODE,
     CONF_NAME,
     CONF_UNIT_SYSTEM,
+    CONF_UNIT_SYSTEM_IMPERIAL,
+    LENGTH_KILOMETERS,
+    LENGTH_MILES,
     TIME_MINUTES,
 )
 from homeassistant.core import HomeAssistant
@@ -139,12 +142,6 @@ def sensor_descriptions(travel_mode: str) -> tuple[SensorEntityDescription, ...]
             state_class=SensorStateClass.MEASUREMENT,
             native_unit_of_measurement=TIME_MINUTES,
         ),
-        SensorEntityDescription(
-            name="Distance",
-            icon=ICONS.get(travel_mode, ICON_CAR),
-            key=ATTR_DISTANCE,
-            state_class=SensorStateClass.MEASUREMENT,
-        ),
         SensorEntityDescription(
             name="Route",
             icon="mdi:directions",
@@ -198,6 +195,7 @@ async def async_setup_entry(
         )
     sensors.append(OriginSensor(entry_id, name, coordinator))
     sensors.append(DestinationSensor(entry_id, name, coordinator))
+    sensors.append(DistanceSensor(entry_id, name, coordinator))
     async_add_entities(sensors)
 
 
@@ -301,3 +299,29 @@ class DestinationSensor(HERETravelTimeSensor):
                 ATTR_LONGITUDE: self.coordinator.data[ATTR_DESTINATION].split(",")[1],
             }
         return None
+
+
+class DistanceSensor(HERETravelTimeSensor):
+    """Sensor holding information about the distance."""
+
+    def __init__(
+        self,
+        unique_id_prefix: str,
+        name: str,
+        coordinator: HereTravelTimeDataUpdateCoordinator,
+    ) -> None:
+        """Initialize the sensor."""
+        sensor_description = SensorEntityDescription(
+            name="Distance",
+            icon=ICONS.get(coordinator.config.travel_mode, ICON_CAR),
+            key=ATTR_DISTANCE,
+            state_class=SensorStateClass.MEASUREMENT,
+        )
+        super().__init__(unique_id_prefix, name, sensor_description, coordinator)
+
+    @property
+    def native_unit_of_measurement(self) -> str | None:
+        """Return the unit of measurement of the sensor."""
+        if self.coordinator.config.units == CONF_UNIT_SYSTEM_IMPERIAL:
+            return LENGTH_MILES
+        return LENGTH_KILOMETERS
diff --git a/tests/components/here_travel_time/test_sensor.py b/tests/components/here_travel_time/test_sensor.py
index 70669a63da5..849c123e72e 100644
--- a/tests/components/here_travel_time/test_sensor.py
+++ b/tests/components/here_travel_time/test_sensor.py
@@ -37,10 +37,13 @@ from homeassistant.const import (
     ATTR_ICON,
     ATTR_LATITUDE,
     ATTR_LONGITUDE,
+    ATTR_UNIT_OF_MEASUREMENT,
     CONF_API_KEY,
     CONF_MODE,
     CONF_NAME,
     EVENT_HOMEASSISTANT_START,
+    LENGTH_KILOMETERS,
+    LENGTH_MILES,
     TIME_MINUTES,
 )
 from homeassistant.core import HomeAssistant
@@ -58,7 +61,7 @@ from tests.common import MockConfigEntry
 
 
 @pytest.mark.parametrize(
-    "mode,icon,unit_system,arrival_time,departure_time,expected_duration,expected_distance,expected_duration_in_traffic",
+    "mode,icon,unit_system,arrival_time,departure_time,expected_duration,expected_distance,expected_duration_in_traffic,expected_distance_unit",
     [
         (
             TRAVEL_MODE_CAR,
@@ -69,6 +72,7 @@ from tests.common import MockConfigEntry
             "30",
             23.903,
             "31",
+            LENGTH_KILOMETERS,
         ),
         (
             TRAVEL_MODE_BICYCLE,
@@ -79,6 +83,7 @@ from tests.common import MockConfigEntry
             "30",
             23.903,
             "30",
+            LENGTH_KILOMETERS,
         ),
         (
             TRAVEL_MODE_PEDESTRIAN,
@@ -89,6 +94,7 @@ from tests.common import MockConfigEntry
             "30",
             14.85263,
             "30",
+            LENGTH_MILES,
         ),
         (
             TRAVEL_MODE_PUBLIC_TIME_TABLE,
@@ -99,6 +105,7 @@ from tests.common import MockConfigEntry
             "30",
             14.85263,
             "30",
+            LENGTH_MILES,
         ),
         (
             TRAVEL_MODE_TRUCK,
@@ -109,6 +116,7 @@ from tests.common import MockConfigEntry
             "30",
             23.903,
             "31",
+            LENGTH_KILOMETERS,
         ),
     ],
 )
@@ -123,6 +131,7 @@ async def test_sensor(
     expected_duration,
     expected_distance,
     expected_duration_in_traffic,
+    expected_distance_unit,
 ):
     """Test that sensor works."""
     entry = MockConfigEntry(
@@ -166,6 +175,10 @@ async def test_sensor(
     assert float(hass.states.get("sensor.test_distance").state) == pytest.approx(
         expected_distance
     )
+    assert (
+        hass.states.get("sensor.test_distance").attributes.get(ATTR_UNIT_OF_MEASUREMENT)
+        == expected_distance_unit
+    )
     assert hass.states.get("sensor.test_route").state == (
         "US-29 - K St NW; US-29 - Whitehurst Fwy; "
         "I-495 N - Capital Beltway; MD-187 S - Old Georgetown Rd"