Fix handling invalid enity_category for sensors

home-assistant · Nov 2, 2023 · 69b0ee4 · 69b0ee4
1 parent 0e21fef
commit 69b0ee4
Show file tree

Hide file tree

Showing 4 changed files with 58 additions and 3 deletions.
diff --git a/homeassistant/components/mqtt/binary_sensor.py b/homeassistant/components/mqtt/binary_sensor.py
@@ -42,6 +42,7 @@
     MqttAvailability,
     MqttEntity,
     async_setup_entity_entry_helper,
+    validate_sensor_entity_category,
     write_state_on_attr_change,
 )
 from .models import MqttValueTemplate, ReceiveMessage
@@ -55,7 +56,7 @@
 DEFAULT_FORCE_UPDATE = False
 CONF_EXPIRE_AFTER = "expire_after"
 
-PLATFORM_SCHEMA_MODERN = MQTT_RO_SCHEMA.extend(
+_PLATFORM_SCHEMA_BASE = MQTT_RO_SCHEMA.extend(
     {
         vol.Optional(CONF_DEVICE_CLASS): vol.Any(DEVICE_CLASSES_SCHEMA, None),
         vol.Optional(CONF_EXPIRE_AFTER): cv.positive_int,
@@ -67,7 +68,12 @@
     }
 ).extend(MQTT_ENTITY_COMMON_SCHEMA.schema)
 
-DISCOVERY_SCHEMA = PLATFORM_SCHEMA_MODERN.extend({}, extra=vol.REMOVE_EXTRA)
+DISCOVERY_SCHEMA = vol.All(
+    validate_sensor_entity_category,
+    _PLATFORM_SCHEMA_BASE.extend({}, extra=vol.REMOVE_EXTRA),
+)
+
+PLATFORM_SCHEMA_MODERN = vol.All(validate_sensor_entity_category, _PLATFORM_SCHEMA_BASE)
 
 
 async def async_setup_entry(

diff --git a/homeassistant/components/mqtt/mixins.py b/homeassistant/components/mqtt/mixins.py
@@ -28,6 +28,7 @@
     CONF_NAME,
     CONF_UNIQUE_ID,
     CONF_VALUE_TEMPLATE,
+    EntityCategory,
 )
 from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback
 from homeassistant.helpers import (
@@ -208,6 +209,16 @@ def validate_device_has_at_least_one_identifier(value: ConfigType) -> ConfigType
     )
 
 
+def validate_sensor_entity_category(config: ConfigType) -> ConfigType:
+    """Check the sensor's entity category is `diagnostic` or `None`."""
+    if (
+        CONF_ENTITY_CATEGORY in config
+        and config[CONF_ENTITY_CATEGORY] == EntityCategory.CONFIG
+    ):
+        raise vol.Invalid("Entity category `config` is invalid")
+    return config
+
+
 MQTT_ENTITY_DEVICE_INFO_SCHEMA = vol.All(
     cv.deprecated(CONF_DEPRECATED_VIA_HUB, CONF_VIA_DEVICE),
     vol.Schema(

diff --git a/homeassistant/components/mqtt/sensor.py b/homeassistant/components/mqtt/sensor.py
@@ -44,6 +44,7 @@
     MqttAvailability,
     MqttEntity,
     async_setup_entity_entry_helper,
+    validate_sensor_entity_category,
     write_state_on_attr_change,
 )
 from .models import (
@@ -70,7 +71,6 @@
 DEFAULT_NAME = "MQTT Sensor"
 DEFAULT_FORCE_UPDATE = False
 
-
 _PLATFORM_SCHEMA_BASE = MQTT_RO_SCHEMA.extend(
     {
         vol.Optional(CONF_DEVICE_CLASS): vol.Any(DEVICE_CLASSES_SCHEMA, None),
@@ -88,13 +88,15 @@
     # Deprecated in HA Core 2021.11.0 https://github.com/home-assistant/core/pull/54840
     # Removed in HA Core 2023.6.0
     cv.removed(CONF_LAST_RESET_TOPIC),
+    validate_sensor_entity_category,
     _PLATFORM_SCHEMA_BASE,
 )
 
 DISCOVERY_SCHEMA = vol.All(
     # Deprecated in HA Core 2021.11.0 https://github.com/home-assistant/core/pull/54840
     # Removed in HA Core 2023.6.0
     cv.removed(CONF_LAST_RESET_TOPIC),
+    validate_sensor_entity_category,
     _PLATFORM_SCHEMA_BASE.extend({}, extra=vol.REMOVE_EXTRA),
 )
 

diff --git a/tests/components/mqtt/test_init.py b/tests/components/mqtt/test_init.py
@@ -2151,6 +2151,42 @@ async def test_setup_manual_mqtt_with_invalid_config(
     assert "required key not provided" in caplog.text
 
 
+@pytest.mark.parametrize(
+    "hass_config",
+    [
+        {
+            mqtt.DOMAIN: {
+                "sensor": {
+                    "name": "test",
+                    "state_topic": "test-topic",
+                    "entity_category": "config",
+                }
+            }
+        },
+        {
+            mqtt.DOMAIN: {
+                "binary_sensor": {
+                    "name": "test",
+                    "state_topic": "test-topic",
+                    "entity_category": "config",
+                }
+            }
+        },
+    ],
+)
+@patch(
+    "homeassistant.components.mqtt.PLATFORMS", [Platform.BINARY_SENSOR, Platform.SENSOR]
+)
+async def test_setup_manual_mqtt_with_invalid_entity_category(
+    hass: HomeAssistant,
+    mqtt_mock_entry: MqttMockHAClientGenerator,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test set up a manual sensor item with an invalid entity category."""
+    assert await mqtt_mock_entry()
+    assert "Entity category `config` is invalid" in caplog.text
+
+
 @patch("homeassistant.components.mqtt.PLATFORMS", [])
 @pytest.mark.parametrize(
     ("mqtt_config_entry_data", "protocol"),