Fix sensor unit conversion bug (#88825)

* Fix sensor unit conversion bug

* Ensure the correct unit is stored in the entity registry

homeassistant/components/sensor/__init__.py

@@ -196,19 +196,30 @@ class SensorEntity(Entity):
         if self.unique_id is None or self.device_class is None:
             return
         registry = er.async_get(self.hass)
+
+        # Bail out if the entity is not yet registered
         if not (
             entity_id := registry.async_get_entity_id(
                 platform.domain, platform.platform_name, self.unique_id
             )
         ):
+            # Prime _sensor_option_unit_of_measurement to ensure the correct unit
+            # is stored in the entity registry.
+            self._sensor_option_unit_of_measurement = self._get_initial_suggested_unit()
             return
+
         registry_entry = registry.async_get(entity_id)
         assert registry_entry
 
+        # Prime _sensor_option_unit_of_measurement to ensure the correct unit
+        # is stored in the entity registry.
+        self.registry_entry = registry_entry
+        self._async_read_entity_options()
+
         # If the sensor has 'unit_of_measurement' in its sensor options, the user has
         # overridden the unit.
-        # If the sensor has 'sensor.private' in its entity options, it was added after
+        # If the sensor has 'sensor.private' in its entity options, it already has a
-        # automatic unit conversion was implemented.
+        # suggested_unit.
         registry_unit = registry_entry.unit_of_measurement
         if (
             (
@@ -230,11 +241,14 @@ class SensorEntity(Entity):
 
         # Set suggested_unit_of_measurement to the old unit to enable automatic
         # conversion
-        registry.async_update_entity_options(
+        self.registry_entry = registry.async_update_entity_options(
             entity_id,
             f"{DOMAIN}.private",
             {"suggested_unit_of_measurement": registry_unit},
         )
+        # Update _sensor_option_unit_of_measurement to ensure the correct unit
+        # is stored in the entity registry.
+        self._async_read_entity_options()
 
     async def async_internal_added_to_hass(self) -> None:
         """Call when the sensor entity is added to hass."""
@@ -305,12 +319,8 @@ class SensorEntity(Entity):
 
         return None
 
-    def get_initial_entity_options(self) -> er.EntityOptionsType | None:
+    def _get_initial_suggested_unit(self) -> str | UndefinedType:
-        """Return initial entity options.
+        """Return the initial unit."""
-
-        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
 
@@ -321,6 +331,19 @@ class SensorEntity(Entity):
             )
 
         if suggested_unit_of_measurement is None:
+            return UNDEFINED
+
+        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 UNDEFINED:
             return None
 
         return {
@@ -416,7 +439,7 @@ class SensorEntity(Entity):
             return self._sensor_option_unit_of_measurement
 
         # Second priority, for non registered entities: unit suggested by integration
-        if not self.registry_entry and self.suggested_unit_of_measurement:
+        if not self.unique_id and self.suggested_unit_of_measurement:
             return self.suggested_unit_of_measurement
 
         # Third priority: Legacy temperature conversion, which applies

tests/components/sensor/test_init.py

@@ -915,6 +915,7 @@ async def test_unit_conversion_priority(
     assert state.attributes[ATTR_UNIT_OF_MEASUREMENT] == automatic_unit
     # Assert the automatic unit conversion is stored in the registry
     entry = entity_registry.async_get(entity0.entity_id)
+    assert entry.unit_of_measurement == automatic_unit
     assert entry.options == {
         "sensor.private": {"suggested_unit_of_measurement": automatic_unit}
     }
@@ -930,6 +931,7 @@ async def test_unit_conversion_priority(
     assert state.attributes[ATTR_UNIT_OF_MEASUREMENT] == suggested_unit
     # Assert the suggested unit is stored in the registry
     entry = entity_registry.async_get(entity2.entity_id)
+    assert entry.unit_of_measurement == suggested_unit
     assert entry.options == {
         "sensor.private": {"suggested_unit_of_measurement": suggested_unit}
     }
@@ -1065,6 +1067,7 @@ async def test_unit_conversion_priority_precision(
     assert state.attributes[ATTR_UNIT_OF_MEASUREMENT] == automatic_unit
     # Assert the automatic unit conversion is stored in the registry
     entry = entity_registry.async_get(entity0.entity_id)
+    assert entry.unit_of_measurement == automatic_unit
     assert entry.options == {
         "sensor": {"suggested_display_precision": 2},
         "sensor.private": {"suggested_unit_of_measurement": automatic_unit},
@@ -1081,6 +1084,7 @@ async def test_unit_conversion_priority_precision(
     assert state.attributes[ATTR_UNIT_OF_MEASUREMENT] == suggested_unit
     # Assert the suggested unit is stored in the registry
     entry = entity_registry.async_get(entity2.entity_id)
+    assert entry.unit_of_measurement == suggested_unit
     assert entry.options == {
         "sensor": {"suggested_display_precision": 2},
         "sensor.private": {"suggested_unit_of_measurement": suggested_unit},
@@ -1154,13 +1158,17 @@ async def test_unit_conversion_priority_suggested_unit_change(
     platform.init(empty=True)
 
     # Pre-register entities
-    entry = entity_registry.async_get_or_create("sensor", "test", "very_unique")
+    entry = entity_registry.async_get_or_create(
+        "sensor", "test", "very_unique", unit_of_measurement=original_unit
+    )
     entity_registry.async_update_entity_options(
         entry.entity_id,
         "sensor.private",
         {"suggested_unit_of_measurement": original_unit},
     )
-    entry = entity_registry.async_get_or_create("sensor", "test", "very_unique_2")
+    entry = entity_registry.async_get_or_create(
+        "sensor", "test", "very_unique_2", unit_of_measurement=original_unit
+    )
     entity_registry.async_update_entity_options(
         entry.entity_id,
         "sensor.private",
@@ -1193,11 +1201,124 @@ async def test_unit_conversion_priority_suggested_unit_change(
     state = hass.states.get(entity0.entity_id)
     assert float(state.state) == pytest.approx(float(original_value))
     assert state.attributes[ATTR_UNIT_OF_MEASUREMENT] == original_unit
+    # Assert the suggested unit is stored in the registry
+    entry = entity_registry.async_get(entity0.entity_id)
+    assert entry.unit_of_measurement == original_unit
+    assert entry.options == {
+        "sensor.private": {"suggested_unit_of_measurement": original_unit},
+    }
 
     # Registered entity -> Follow suggested unit the first time the entity was seen
     state = hass.states.get(entity1.entity_id)
     assert float(state.state) == pytest.approx(float(original_value))
     assert state.attributes[ATTR_UNIT_OF_MEASUREMENT] == original_unit
+    # Assert the suggested unit is stored in the registry
+    entry = entity_registry.async_get(entity1.entity_id)
+    assert entry.unit_of_measurement == original_unit
+    assert entry.options == {
+        "sensor.private": {"suggested_unit_of_measurement": original_unit},
+    }
+
+
+@pytest.mark.parametrize(
+    (
+        "native_unit_1",
+        "native_unit_2",
+        "suggested_unit",
+        "native_value",
+        "original_value",
+        "device_class",
+    ),
+    [
+        # Distance
+        (
+            UnitOfLength.KILOMETERS,
+            UnitOfLength.METERS,
+            UnitOfLength.KILOMETERS,
+            1000000,
+            1000,
+            SensorDeviceClass.DISTANCE,
+        ),
+        # Energy
+        (
+            UnitOfEnergy.KILO_WATT_HOUR,
+            UnitOfEnergy.WATT_HOUR,
+            UnitOfEnergy.KILO_WATT_HOUR,
+            1000000,
+            1000,
+            SensorDeviceClass.ENERGY,
+        ),
+    ],
+)
+async def test_unit_conversion_priority_suggested_unit_change_2(
+    hass: HomeAssistant,
+    enable_custom_integrations: None,
+    native_unit_1,
+    native_unit_2,
+    suggested_unit,
+    native_value,
+    original_value,
+    device_class,
+) -> None:
+    """Test priority of unit conversion."""
+
+    hass.config.units = METRIC_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=native_unit_1
+    )
+    entity_registry.async_get_or_create(
+        "sensor", "test", "very_unique_2", unit_of_measurement=native_unit_1
+    )
+
+    platform.ENTITIES["0"] = platform.MockSensor(
+        name="Test",
+        device_class=device_class,
+        native_unit_of_measurement=native_unit_2,
+        native_value=str(native_value),
+        unique_id="very_unique",
+    )
+    entity0 = platform.ENTITIES["0"]
+
+    platform.ENTITIES["1"] = platform.MockSensor(
+        name="Test",
+        device_class=device_class,
+        native_unit_of_measurement=native_unit_2,
+        native_value=str(native_value),
+        suggested_unit_of_measurement=suggested_unit,
+        unique_id="very_unique_2",
+    )
+    entity1 = platform.ENTITIES["1"]
+
+    assert await async_setup_component(hass, "sensor", {"sensor": {"platform": "test"}})
+    await hass.async_block_till_done()
+
+    # Registered entity -> Follow unit in entity registry
+    state = hass.states.get(entity0.entity_id)
+    assert float(state.state) == pytest.approx(float(original_value))
+    assert state.attributes[ATTR_UNIT_OF_MEASUREMENT] == native_unit_1
+    # Assert the suggested unit is stored in the registry
+    entry = entity_registry.async_get(entity0.entity_id)
+    assert entry.unit_of_measurement == native_unit_1
+    assert entry.options == {
+        "sensor.private": {"suggested_unit_of_measurement": native_unit_1},
+    }
+
+    # Registered entity -> Follow unit in entity registry
+    state = hass.states.get(entity1.entity_id)
+    assert float(state.state) == pytest.approx(float(original_value))
+    assert state.attributes[ATTR_UNIT_OF_MEASUREMENT] == native_unit_1
+    # Assert the suggested unit is stored in the registry
+    entry = entity_registry.async_get(entity0.entity_id)
+    assert entry.unit_of_measurement == native_unit_1
+    assert entry.options == {
+        "sensor.private": {"suggested_unit_of_measurement": native_unit_1},
+    }
 
 
 @pytest.mark.parametrize(