Fix sensor unit conversion bug (#88825)

* Fix sensor unit conversion bug

* Ensure the correct unit is stored in the entity registry
pull/88979/head
Erik Montnemery 2023-02-27 11:46:55 +01:00 committed by Paulus Schoutsen
parent b054c81e13
commit 00b59c142a
2 changed files with 156 additions and 12 deletions

View File

@ -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
# automatic unit conversion was implemented.
# If the sensor has 'sensor.private' in its entity options, it already has a
# 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:
"""Return initial entity options.
These will be stored in the entity registry the first time the entity is seen,
and then never updated.
"""
def _get_initial_suggested_unit(self) -> str | UndefinedType:
"""Return the initial unit."""
# 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

View File

@ -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(