Imprve handling of native_unit_of_measurement

Author: tomer-w

custom_components/victron_mqtt/entity.py

@@ -69,7 +69,9 @@ def __init__(
             metric.generic_short_id not in ENTITIES_DISABLE_BY_DEFAULT
         )
 
-    def _native_unit_of_measurement(self) -> str | None:
+    @property
+    def native_unit_of_measurement(self) -> str | None:
+        """Return the metric unit before entity registration."""
         if self._metric.metric_type == MetricType.COST:
             return self.hass.config.currency
 

custom_components/victron_mqtt/number.py

@@ -92,7 +92,6 @@ def __init__(
             simple_naming,
             installation_id,
         )
-        self._attr_native_unit_of_measurement = None
         self._attr_device_class = METRIC_TYPE_TO_DEVICE_CLASS.get(metric.metric_type)
         self._attr_native_value = metric.value
         if metric.min_value is not None:
@@ -107,11 +106,6 @@ def _on_update_cb(self, value: Any) -> None:
         self._attr_native_value = value
         self.async_write_ha_state()
 
-    async def async_added_to_hass(self) -> None:
-        """Run when entity about to be added to hass."""
-        self._attr_native_unit_of_measurement = self._native_unit_of_measurement()
-        await super().async_added_to_hass()
-
     async def async_set_native_value(self, value: float) -> None:
         """Set a new value."""
         if TYPE_CHECKING:

custom_components/victron_mqtt/sensor.py

@@ -105,7 +105,6 @@ def __init__(
         super().__init__(
             device, metric, device_info, "sensor", simple_naming, installation_id
         )
-        self._attr_native_unit_of_measurement = None
         self._attr_device_class = METRIC_TYPE_TO_DEVICE_CLASS.get(metric.metric_type)
         # Enum sensors must not have a state class
         if self._attr_device_class == SensorDeviceClass.ENUM:
@@ -133,8 +132,6 @@ def _normalize_value(value: Any) -> Any:
 
     async def async_added_to_hass(self) -> None:
         """Restore persistent state for FormulaMetric energy sensors."""
-        self._attr_native_unit_of_measurement = self._native_unit_of_measurement()
-
         # Only restore for cumulative FormulaMetric sensors (TOTAL / TOTAL_INCREASING).
         # These metrics start from 0 on each HA restart, so we restore the
         # previous accumulated value as a baseline and add new increments on top.

tests/test_hub.py

@@ -323,6 +323,27 @@ async def test_sensor_without_unit_does_not_crash(
     assert state.attributes.get("unit_of_measurement") is None
 
 
+async def test_sensor_uses_registry_unit(
+    hass: HomeAssistant,
+    init_integration,
+) -> None:
+    """Test sensor unit is available when the entity is registered."""
+    victron_hub, mock_config_entry = init_integration
+
+    await inject_message(victron_hub, "N/123/battery/0/Dc/0/Current", '{"value": 10.5}')
+    await finalize_injection(victron_hub)
+    await hass.async_block_till_done()
+
+    entity_registry = er.async_get(hass)
+    entities = er.async_entries_for_config_entry(
+        entity_registry, mock_config_entry.entry_id
+    )
+    assert len(entities) > 0
+
+    entity = next(entry for entry in entities if entry.translation_key == "battery_current")
+    assert entity.unit_of_measurement == "A"
+
+
 async def test_monetary_sensor_uses_ha_currency(
     hass: HomeAssistant,
     init_integration,
@@ -340,6 +361,7 @@ async def test_monetary_sensor_uses_ha_currency(
         entity_registry, mock_config_entry.entry_id
     )
     assert len(entities) == 1
+    assert entities[0].unit_of_measurement == "EUR"
 
     state = hass.states.get(entities[0].entity_id)
     assert state is not None
@@ -1091,4 +1113,3 @@ async def test_number_with_step(
     assert state is not None
     assert float(state.state) == 57.6
     assert state.attributes.get("step") == 0.1
-