Allow integrations to drop custom unit conversion (#81005)

Co-authored-by: Paulus Schoutsen <balloob@gmail.com>
Co-authored-by: Franck Nijhof <git@frenck.dev>

homeassistant/components/sensor/__init__.py

@@ -1,6 +1,7 @@
 """Component to interface with various sensors that can be monitored."""
 from __future__ import annotations
 
+import asyncio
 from collections.abc import Mapping
 from contextlib import suppress
 from dataclasses import dataclass
@@ -56,6 +57,7 @@ from homeassistant.helpers.config_validation import (  # noqa: F401
 )
 from homeassistant.helpers.entity import Entity, EntityDescription
 from homeassistant.helpers.entity_component import EntityComponent
+from homeassistant.helpers.entity_platform import EntityPlatform
 from homeassistant.helpers.restore_state import ExtraStoredData, RestoreEntity
 from homeassistant.helpers.typing import ConfigType, StateType
 from homeassistant.util import dt as dt_util
@@ -453,6 +455,47 @@ class SensorEntity(Entity):
     _last_reset_reported = False
     _sensor_option_unit_of_measurement: str | None = None
 
+    @callback
+    def add_to_platform_start(
+        self,
+        hass: HomeAssistant,
+        platform: EntityPlatform,
+        parallel_updates: asyncio.Semaphore | None,
+    ) -> None:
+        """Start adding an entity to a platform."""
+        super().add_to_platform_start(hass, platform, parallel_updates)
+
+        if self.unique_id is None:
+            return
+        registry = er.async_get(self.hass)
+        if not (
+            entity_id := registry.async_get_entity_id(
+                platform.domain, platform.platform_name, self.unique_id
+            )
+        ):
+            return
+        registry_entry = registry.async_get(entity_id)
+        assert registry_entry
+
+        # Store unit override according to automatic unit conversion rules if:
+        # - no unit override is stored in the entity registry
+        # - units have changed
+        # - the unit stored in the registry matches automatic unit conversion rules
+        # This allows integrations to drop custom unit conversion and rely on automatic
+        # conversion.
+        registry_unit = registry_entry.unit_of_measurement
+        if (
+            DOMAIN not in registry_entry.options
+            and f"{DOMAIN}.private" not in registry_entry.options
+            and self.unit_of_measurement != registry_unit
+            and (suggested_unit := self._get_initial_suggested_unit()) == registry_unit
+        ):
+            registry.async_update_entity_options(
+                entity_id,
+                f"{DOMAIN}.private",
+                {"suggested_unit_of_measurement": suggested_unit},
+            )
+
     async def async_internal_added_to_hass(self) -> None:
         """Call when the sensor entity is added to hass."""
         await super().async_internal_added_to_hass()
@@ -495,12 +538,8 @@ class SensorEntity(Entity):
 
         return None
 
-    def get_initial_entity_options(self) -> er.EntityOptionsType | None:
+    def _get_initial_suggested_unit(self) -> str | None:
-        """Return initial entity options.
+        """Return initial suggested unit of measurement."""
-
-        These will be stored in the entity registry the first time the entity is seen,
-        and then never updated.
-        """
         # Unit suggested by the integration
         suggested_unit_of_measurement = self.suggested_unit_of_measurement
 
@@ -510,6 +549,15 @@ class SensorEntity(Entity):
                 self.device_class, self.native_unit_of_measurement
             )
 
+        return suggested_unit_of_measurement
+
+    def get_initial_entity_options(self) -> er.EntityOptionsType | None:
+        """Return initial entity options.
+
+        These will be stored in the entity registry the first time the entity is seen,
+        and then never updated.
+        """
+        suggested_unit_of_measurement = self._get_initial_suggested_unit()
         if suggested_unit_of_measurement is None:
             return None
 

tests/components/sensor/test_init.py

@@ -873,3 +873,57 @@ async def test_unit_conversion_priority_suggested_unit_change(
     state = hass.states.get(entity1.entity_id)
     assert float(state.state) == approx(float(original_value))
     assert state.attributes[ATTR_UNIT_OF_MEASUREMENT] == original_unit
+
+
+@pytest.mark.parametrize(
+    "unit_system, native_unit, original_unit, native_value, original_value, device_class",
+    [
+        # Distance
+        (
+            US_CUSTOMARY_SYSTEM,
+            LENGTH_KILOMETERS,
+            LENGTH_MILES,
+            1000,
+            621,
+            SensorDeviceClass.DISTANCE,
+        ),
+    ],
+)
+async def test_unit_conversion_priority_legacy_conversion_removed(
+    hass,
+    enable_custom_integrations,
+    unit_system,
+    native_unit,
+    original_unit,
+    native_value,
+    original_value,
+    device_class,
+):
+    """Test priority of unit conversion."""
+
+    hass.config.units = unit_system
+
+    entity_registry = er.async_get(hass)
+    platform = getattr(hass.components, "test.sensor")
+    platform.init(empty=True)
+
+    # Pre-register entities
+    entity_registry.async_get_or_create(
+        "sensor", "test", "very_unique", unit_of_measurement=original_unit
+    )
+
+    platform.ENTITIES["0"] = platform.MockSensor(
+        name="Test",
+        device_class=device_class,
+        native_unit_of_measurement=native_unit,
+        native_value=str(native_value),
+        unique_id="very_unique",
+    )
+    entity0 = platform.ENTITIES["0"]
+
+    assert await async_setup_component(hass, "sensor", {"sensor": {"platform": "test"}})
+    await hass.async_block_till_done()
+
+    state = hass.states.get(entity0.entity_id)
+    assert float(state.state) == approx(float(original_value))
+    assert state.attributes[ATTR_UNIT_OF_MEASUREMENT] == original_unit