Add validation on unit of measurement and show SelectSelector if we know the device class

Author: jbouwh

homeassistant/components/mqtt/config_flow.py

@@ -29,6 +29,7 @@
 from homeassistant.components.hassio import AddonError, AddonManager, AddonState
 from homeassistant.components.sensor import (
     CONF_STATE_CLASS,
+    DEVICE_CLASS_UNITS,
     SensorDeviceClass,
     SensorStateClass,
 )
@@ -271,10 +272,28 @@
 class PlatformField:
     """Stores a platform config field schema, required flag and validator."""
 
-    selector: Selector
+    selector: Selector[Any] | Callable[..., Selector[Any]]
     required: bool
     validator: Callable[..., Any]
     error: str | None = None
+    custom_filtering: bool = False
+
+
+@callback
+def unit_of_measurement_selector(user_data: dict[str, Any | None]) -> Selector:
+    """Return a context based unit of measurement selector."""
+    if (
+        user_data is None
+        or (device_class := user_data.get(CONF_DEVICE_CLASS)) is None
+        or device_class not in DEVICE_CLASS_UNITS
+    ):
+        return TEXT_SELECTOR
+    return SelectSelector(
+        SelectSelectorConfig(
+            options=[str(uom) for uom in DEVICE_CLASS_UNITS[device_class]],
+            custom_value=True,
+        )
+    )
 
 
 COMMON_ENTITY_FIELDS = {
@@ -291,7 +310,9 @@ class PlatformField:
     Platform.SENSOR.value: {
         CONF_DEVICE_CLASS: PlatformField(SENSOR_DEVICE_CLASS_SELECTOR, False, str),
         CONF_STATE_CLASS: PlatformField(SENSOR_STATE_CLASS_SELECTOR, False, str),
-        CONF_UNIT_OF_MEASUREMENT: PlatformField(TEXT_SELECTOR, False, str),
+        CONF_UNIT_OF_MEASUREMENT: PlatformField(
+            unit_of_measurement_selector, False, str, custom_filtering=True
+        ),
         CONF_SUGGESTED_DISPLAY_PRECISION: PlatformField(
             SUGGESTED_DISPLAY_PRECISION_SELECTOR, False, cv.positive_int
         ),
@@ -413,13 +434,22 @@ def validate_user_input(
 
 
 @callback
-def data_schema_from_fields(data_schema_fields: dict[str, PlatformField]) -> vol.Schema:
-    """Generate data schema from platform fields."""
+def data_schema_from_fields(
+    data_schema_fields: dict[str, PlatformField],
+    component: dict[str, Any] | None = None,
+    user_input: dict[str, Any] | None = None,
+) -> vol.Schema:
+    """Generate custom data schema from platform fields."""
+    user_data = component
+    if user_data is not None and user_input is not None:
+        user_data |= user_input
     return vol.Schema(
         {
             vol.Required(field_name)
             if field_details.required
-            else vol.Optional(field_name): field_details.selector
+            else vol.Optional(field_name): field_details.selector(user_data)  # type: ignore[operator]
+            if field_details.custom_filtering
+            else field_details.selector
             for field_name, field_details in data_schema_fields.items()
         }
     )
@@ -1141,13 +1171,14 @@ async def async_step_entity_platform_config(
         """Configure platform entity details."""
         if TYPE_CHECKING:
             assert self._component_id is not None
-        platform = self._subentry_data["components"][self._component_id][CONF_PLATFORM]
+        component = self._subentry_data["components"][self._component_id]
+        platform = component[CONF_PLATFORM]
         if not (data_schema_fields := PLATFORM_ENTITY_FIELDS[platform]):
             return await self.async_step_mqtt_platform_config()
         errors: dict[str, str] = {}
         device_name, full_entity_name = self.generate_names()
 
-        data_schema = data_schema_from_fields(data_schema_fields)
+        data_schema = data_schema_from_fields(data_schema_fields, component, user_input)
         if user_input is not None:
             # Test entity fields against the validator
             self.reset_if_empty(user_input)

homeassistant/components/mqtt/strings.json

@@ -251,9 +251,10 @@
         "invalid_input": "Invalid value",
         "invalid_subscribe_topic": "Invalid subscribe topic",
         "invalid_template": "Invalid template",
+        "invalid_uom": "The unit of measurement is not allowed with the selected device class, please use the correct device class, or pick a valid unit of measurement from the list",
         "invalid_url": "Invalid URL",
         "last_reset_not_with_state_class_total": "The last reset value template option should be used with state class 'Total' only",
-        "options_not_allowed_with_state_class_or_uom": "The 'Options' setting is not allowed when state class or unit of measurement are used.",
+        "options_not_allowed_with_state_class_or_uom": "The 'Options' setting is not allowed when state class or unit of measurement are used",
         "options_device_class_enum": "The 'Options' setting must be used with the Enumeration device class'"
       }
     }

homeassistant/components/mqtt/util.py

@@ -15,6 +15,7 @@
 
 from homeassistant.components.sensor import (
     CONF_STATE_CLASS,
+    DEVICE_CLASS_UNITS,
     SensorDeviceClass,
     SensorStateClass,
 )
@@ -426,11 +427,13 @@ def migrate_certificate_file_to_content(file_name_or_auto: str) -> str | None:
         return None
 
 
+@callback
 def learn_more_url(platform: str) -> str:
     """Return the URL for the platform specific MQTT documentation."""
     return f"https://www.home-assistant.io/integrations/{platform}.mqtt/"
 
 
+@callback
 def validate_sensor_state_and_device_class_config(
     config: ConfigType, errors: dict[str, str] | None = None
 ) -> ConfigType:
@@ -468,4 +471,20 @@ def validate_sensor_state_and_device_class_config(
                 )
             errors[CONF_OPTIONS] = "options_device_class_enum"
 
+    if (device_class := config.get(CONF_DEVICE_CLASS)) is None or (
+        unit_of_measurement := config.get(CONF_UNIT_OF_MEASUREMENT)
+    ) is None:
+        return config
+
+    if (
+        device_class in DEVICE_CLASS_UNITS
+        and unit_of_measurement not in DEVICE_CLASS_UNITS[device_class]
+    ):
+        if errors is None:
+            raise vol.Invalid(
+                f"The unit of measurement `{unit_of_measurement}` is not valid "
+                f"together with device class `{device_class}`"
+            )
+        errors[CONF_UNIT_OF_MEASUREMENT] = "invalid_uom"
+
     return config

tests/components/mqtt/test_config_flow.py

@@ -2678,6 +2678,13 @@ async def test_migrate_of_incompatible_config_entry(
                     },
                     {"options": "options_not_allowed_with_state_class_or_uom"},
                 ),
+                (
+                    {
+                        "device_class": "energy",
+                        "unit_of_measurement": "ppm",
+                    },
+                    {"unit_of_measurement": "invalid_uom"},
+                ),
             ),
             {
                 "state_topic": "test-topic",

tests/components/mqtt/test_sensor.py

@@ -870,6 +870,32 @@ async def test_invalid_device_class(
     assert "expected SensorDeviceClass or one of" in caplog.text
 
 
+@pytest.mark.parametrize(
+    "hass_config",
+    [
+        {
+            mqtt.DOMAIN: {
+                sensor.DOMAIN: {
+                    "name": "test",
+                    "state_topic": "test-topic",
+                    "device_class": "energy",
+                    "unit_of_measurement": "ppm",
+                }
+            }
+        }
+    ],
+)
+async def test_invalid_unit_of_measurement(
+    mqtt_mock_entry: MqttMockHAClientGenerator, caplog: pytest.LogCaptureFixture
+) -> None:
+    """Test device_class with invalid unit of measurement."""
+    assert await mqtt_mock_entry()
+    assert (
+        "The unit of measurement `ppm` is not valid together with device class `energy`"
+        in caplog.text
+    )
+
+
 @pytest.mark.parametrize(
     "hass_config",
     [