From 4aa7a9faee0256f76ffa8d9ae0e354860d8c713d Mon Sep 17 00:00:00 2001
From: Thomas Kistler <thomas@uyei-kistler.com>
Date: Fri, 21 Jun 2024 00:07:14 -0700
Subject: [PATCH] Fix Hydrawise volume unit bug (#119988)

---
 homeassistant/components/hydrawise/sensor.py | 15 +++++---
 tests/components/hydrawise/conftest.py       |  7 +++-
 tests/components/hydrawise/test_sensor.py    | 36 ++++++++++++++++++--
 3 files changed, 51 insertions(+), 7 deletions(-)

diff --git a/homeassistant/components/hydrawise/sensor.py b/homeassistant/components/hydrawise/sensor.py
index 2497fe8f49dd9..fe4b33d58514d 100644
--- a/homeassistant/components/hydrawise/sensor.py
+++ b/homeassistant/components/hydrawise/sensor.py
@@ -71,7 +71,6 @@ def _get_controller_daily_total_water_use(sensor: HydrawiseSensor) -> float | No
         key="daily_total_water_use",
         translation_key="daily_total_water_use",
         device_class=SensorDeviceClass.VOLUME,
-        native_unit_of_measurement=UnitOfVolume.GALLONS,
         suggested_display_precision=1,
         value_fn=_get_controller_daily_total_water_use,
     ),
@@ -79,7 +78,6 @@ def _get_controller_daily_total_water_use(sensor: HydrawiseSensor) -> float | No
         key="daily_active_water_use",
         translation_key="daily_active_water_use",
         device_class=SensorDeviceClass.VOLUME,
-        native_unit_of_measurement=UnitOfVolume.GALLONS,
         suggested_display_precision=1,
         value_fn=_get_controller_daily_active_water_use,
     ),
@@ -87,7 +85,6 @@ def _get_controller_daily_total_water_use(sensor: HydrawiseSensor) -> float | No
         key="daily_inactive_water_use",
         translation_key="daily_inactive_water_use",
         device_class=SensorDeviceClass.VOLUME,
-        native_unit_of_measurement=UnitOfVolume.GALLONS,
         suggested_display_precision=1,
         value_fn=_get_controller_daily_inactive_water_use,
     ),
@@ -98,7 +95,6 @@ def _get_controller_daily_total_water_use(sensor: HydrawiseSensor) -> float | No
         key="daily_active_water_use",
         translation_key="daily_active_water_use",
         device_class=SensorDeviceClass.VOLUME,
-        native_unit_of_measurement=UnitOfVolume.GALLONS,
         suggested_display_precision=1,
         value_fn=_get_zone_daily_active_water_use,
     ),
@@ -165,6 +161,17 @@ class HydrawiseSensor(HydrawiseEntity, SensorEntity):
 
     entity_description: HydrawiseSensorEntityDescription
 
+    @property
+    def native_unit_of_measurement(self) -> str | None:
+        """Return the unit_of_measurement of the sensor."""
+        if self.entity_description.device_class != SensorDeviceClass.VOLUME:
+            return self.entity_description.native_unit_of_measurement
+        return (
+            UnitOfVolume.GALLONS
+            if self.coordinator.data.user.units.units_name == "imperial"
+            else UnitOfVolume.LITERS
+        )
+
     @property
     def icon(self) -> str | None:
         """Icon of the entity based on the value."""
diff --git a/tests/components/hydrawise/conftest.py b/tests/components/hydrawise/conftest.py
index 8bca1de5fed51..eb1518eb7f2e6 100644
--- a/tests/components/hydrawise/conftest.py
+++ b/tests/components/hydrawise/conftest.py
@@ -15,6 +15,7 @@
     Sensor,
     SensorModel,
     SensorStatus,
+    UnitsSummary,
     User,
     Zone,
 )
@@ -85,7 +86,11 @@ def mock_auth() -> Generator[AsyncMock]:
 @pytest.fixture
 def user() -> User:
     """Hydrawise User fixture."""
-    return User(customer_id=12345, email="asdf@asdf.com")
+    return User(
+        customer_id=12345,
+        email="asdf@asdf.com",
+        units=UnitsSummary(units_name="imperial"),
+    )
 
 
 @pytest.fixture
diff --git a/tests/components/hydrawise/test_sensor.py b/tests/components/hydrawise/test_sensor.py
index fcbc47c41f434..af75ad69ade8a 100644
--- a/tests/components/hydrawise/test_sensor.py
+++ b/tests/components/hydrawise/test_sensor.py
@@ -3,13 +3,18 @@
 from collections.abc import Awaitable, Callable
 from unittest.mock import patch
 
-from pydrawise.schema import Controller, Zone
+from pydrawise.schema import Controller, User, Zone
 import pytest
 from syrupy.assertion import SnapshotAssertion
 
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
+from homeassistant.util.unit_system import (
+    METRIC_SYSTEM,
+    US_CUSTOMARY_SYSTEM,
+    UnitSystem,
+)
 
 from tests.common import MockConfigEntry, snapshot_platform
 
@@ -45,7 +50,7 @@ async def test_suspended_state(
     assert next_cycle.state == "unknown"
 
 
-async def test_no_sensor_and_water_state2(
+async def test_no_sensor_and_water_state(
     hass: HomeAssistant,
     controller: Controller,
     mock_add_config_entry: Callable[[], Awaitable[MockConfigEntry]],
@@ -63,3 +68,30 @@ async def test_no_sensor_and_water_state2(
     sensor = hass.states.get("binary_sensor.home_controller_connectivity")
     assert sensor is not None
     assert sensor.state == "on"
+
+
+@pytest.mark.parametrize(
+    ("hydrawise_unit_system", "unit_system", "expected_state"),
+    [
+        ("imperial", METRIC_SYSTEM, "454.6279552584"),
+        ("imperial", US_CUSTOMARY_SYSTEM, "120.1"),
+        ("metric", METRIC_SYSTEM, "120.1"),
+        ("metric", US_CUSTOMARY_SYSTEM, "31.7270634882136"),
+    ],
+)
+async def test_volume_unit_conversion(
+    hass: HomeAssistant,
+    unit_system: UnitSystem,
+    hydrawise_unit_system: str,
+    expected_state: str,
+    user: User,
+    mock_add_config_entry: Callable[[], Awaitable[MockConfigEntry]],
+) -> None:
+    """Test volume unit conversion."""
+    hass.config.units = unit_system
+    user.units.units_name = hydrawise_unit_system
+    await mock_add_config_entry()
+
+    daily_active_water_use = hass.states.get("sensor.zone_one_daily_active_water_use")
+    assert daily_active_water_use is not None
+    assert daily_active_water_use.state == expected_state