From ab0247d1129c62c208686230b3136e55ec8ca622 Mon Sep 17 00:00:00 2001 From: Paulus Schoutsen Date: Tue, 19 Oct 2021 12:29:22 -0700 Subject: [PATCH] Add entity category and state class to mobile app (#58012) --- homeassistant/components/mobile_app/const.py | 2 + homeassistant/components/mobile_app/entity.py | 6 +++ homeassistant/components/mobile_app/sensor.py | 8 +++ .../components/mobile_app/webhook.py | 52 +++++++++++++------ tests/components/mobile_app/test_sensor.py | 10 +++- tests/components/mobile_app/test_webhook.py | 41 +++++++++++++++ 6 files changed, 103 insertions(+), 16 deletions(-) diff --git a/homeassistant/components/mobile_app/const.py b/homeassistant/components/mobile_app/const.py index 48fed416c33..2f5db21b815 100644 --- a/homeassistant/components/mobile_app/const.py +++ b/homeassistant/components/mobile_app/const.py @@ -61,9 +61,11 @@ ERR_INVALID_FORMAT = "invalid_format" ATTR_SENSOR_ATTRIBUTES = "attributes" ATTR_SENSOR_DEVICE_CLASS = "device_class" +ATTR_SENSOR_ENTITY_CATEGORY = "entity_category" ATTR_SENSOR_ICON = "icon" ATTR_SENSOR_NAME = "name" ATTR_SENSOR_STATE = "state" +ATTR_SENSOR_STATE_CLASS = "state_class" ATTR_SENSOR_TYPE = "type" ATTR_SENSOR_TYPE_BINARY_SENSOR = "binary_sensor" ATTR_SENSOR_TYPE_SENSOR = "sensor" diff --git a/homeassistant/components/mobile_app/entity.py b/homeassistant/components/mobile_app/entity.py index d0de89a94a1..0cdec984f55 100644 --- a/homeassistant/components/mobile_app/entity.py +++ b/homeassistant/components/mobile_app/entity.py @@ -9,6 +9,7 @@ from homeassistant.helpers.restore_state import RestoreEntity from .const import ( ATTR_SENSOR_ATTRIBUTES, ATTR_SENSOR_DEVICE_CLASS, + ATTR_SENSOR_ENTITY_CATEGORY, ATTR_SENSOR_ICON, ATTR_SENSOR_STATE, ATTR_SENSOR_TYPE, @@ -86,6 +87,11 @@ class MobileAppEntity(RestoreEntity): """Return the icon to use in the frontend, if any.""" return self._config[ATTR_SENSOR_ICON] + @property + def entity_category(self): + """Return the entity category, if any.""" + return self._config.get(ATTR_SENSOR_ENTITY_CATEGORY) + @property def unique_id(self): """Return the unique ID of this sensor.""" diff --git a/homeassistant/components/mobile_app/sensor.py b/homeassistant/components/mobile_app/sensor.py index f6652f7f889..9d56e55a106 100644 --- a/homeassistant/components/mobile_app/sensor.py +++ b/homeassistant/components/mobile_app/sensor.py @@ -1,4 +1,6 @@ """Sensor platform for mobile_app.""" +from __future__ import annotations + from homeassistant.components.sensor import SensorEntity from homeassistant.const import CONF_NAME, CONF_UNIQUE_ID, CONF_WEBHOOK_ID from homeassistant.core import callback @@ -12,6 +14,7 @@ from .const import ( ATTR_SENSOR_ICON, ATTR_SENSOR_NAME, ATTR_SENSOR_STATE, + ATTR_SENSOR_STATE_CLASS, ATTR_SENSOR_TYPE, ATTR_SENSOR_TYPE_SENSOR as ENTITY_TYPE, ATTR_SENSOR_UNIQUE_ID, @@ -82,3 +85,8 @@ class MobileAppSensor(MobileAppEntity, SensorEntity): def native_unit_of_measurement(self): """Return the unit of measurement this sensor expresses itself in.""" return self._config.get(ATTR_SENSOR_UOM) + + @property + def state_class(self) -> str | None: + """Return state class.""" + return self._config.get(ATTR_SENSOR_STATE_CLASS) diff --git a/homeassistant/components/mobile_app/webhook.py b/homeassistant/components/mobile_app/webhook.py index eb2d64114b3..8a06b693e9a 100644 --- a/homeassistant/components/mobile_app/webhook.py +++ b/homeassistant/components/mobile_app/webhook.py @@ -22,7 +22,10 @@ from homeassistant.components.device_tracker import ( ATTR_LOCATION_NAME, ) from homeassistant.components.frontend import MANIFEST_JSON -from homeassistant.components.sensor import DEVICE_CLASSES as SENSOR_CLASSES +from homeassistant.components.sensor import ( + DEVICE_CLASSES as SENSOR_CLASSES, + STATE_CLASSES as SENSOSR_STATE_CLASSES, +) from homeassistant.components.zone.const import DOMAIN as ZONE_DOMAIN from homeassistant.const import ( ATTR_DEVICE_ID, @@ -41,6 +44,7 @@ from homeassistant.helpers import ( template, ) from homeassistant.helpers.dispatcher import async_dispatcher_send +from homeassistant.helpers.entity import ENTITY_CATEGORIES_SCHEMA from homeassistant.util.decorator import Registry from .const import ( @@ -57,9 +61,11 @@ from .const import ( ATTR_OS_VERSION, ATTR_SENSOR_ATTRIBUTES, ATTR_SENSOR_DEVICE_CLASS, + ATTR_SENSOR_ENTITY_CATEGORY, ATTR_SENSOR_ICON, ATTR_SENSOR_NAME, ATTR_SENSOR_STATE, + ATTR_SENSOR_STATE_CLASS, ATTR_SENSOR_TYPE, ATTR_SENSOR_TYPE_BINARY_SENSOR, ATTR_SENSOR_TYPE_SENSOR, @@ -389,22 +395,38 @@ async def webhook_enable_encryption(hass, config_entry, data): return json_response({"secret": secret}) +def _validate_state_class_sensor(value: dict): + """Validate we only set state class for sensors.""" + if ( + ATTR_SENSOR_STATE_CLASS in value + and value[ATTR_SENSOR_TYPE] != ATTR_SENSOR_TYPE_SENSOR + ): + raise vol.Invalid("state_class only allowed for sensors") + + return value + + @WEBHOOK_COMMANDS.register("register_sensor") @validate_schema( - { - vol.Optional(ATTR_SENSOR_ATTRIBUTES, default={}): dict, - vol.Optional(ATTR_SENSOR_DEVICE_CLASS): vol.All( - vol.Lower, vol.In(COMBINED_CLASSES) - ), - vol.Required(ATTR_SENSOR_NAME): cv.string, - vol.Required(ATTR_SENSOR_TYPE): vol.In(SENSOR_TYPES), - vol.Required(ATTR_SENSOR_UNIQUE_ID): cv.string, - vol.Optional(ATTR_SENSOR_UOM): cv.string, - vol.Optional(ATTR_SENSOR_STATE, default=None): vol.Any( - None, bool, str, int, float - ), - vol.Optional(ATTR_SENSOR_ICON, default="mdi:cellphone"): cv.icon, - } + vol.All( + { + vol.Optional(ATTR_SENSOR_ATTRIBUTES, default={}): dict, + vol.Optional(ATTR_SENSOR_DEVICE_CLASS): vol.All( + vol.Lower, vol.In(COMBINED_CLASSES) + ), + vol.Required(ATTR_SENSOR_NAME): cv.string, + vol.Required(ATTR_SENSOR_TYPE): vol.In(SENSOR_TYPES), + vol.Required(ATTR_SENSOR_UNIQUE_ID): cv.string, + vol.Optional(ATTR_SENSOR_UOM): cv.string, + vol.Optional(ATTR_SENSOR_STATE, default=None): vol.Any( + None, bool, str, int, float + ), + vol.Optional(ATTR_SENSOR_ENTITY_CATEGORY): ENTITY_CATEGORIES_SCHEMA, + vol.Optional(ATTR_SENSOR_ICON, default="mdi:cellphone"): cv.icon, + vol.Optional(ATTR_SENSOR_STATE_CLASS): vol.In(SENSOSR_STATE_CLASSES), + }, + _validate_state_class_sensor, + ) ) async def webhook_register_sensor(hass, config_entry, data): """Handle a register sensor webhook.""" diff --git a/tests/components/mobile_app/test_sensor.py b/tests/components/mobile_app/test_sensor.py index ed638301bd6..fea43ffba9e 100644 --- a/tests/components/mobile_app/test_sensor.py +++ b/tests/components/mobile_app/test_sensor.py @@ -1,6 +1,6 @@ """Entity tests for mobile_app.""" from homeassistant.const import PERCENTAGE, STATE_UNKNOWN -from homeassistant.helpers import device_registry as dr +from homeassistant.helpers import device_registry as dr, entity_registry as er async def test_sensor(hass, create_registrations, webhook_client): @@ -19,7 +19,9 @@ async def test_sensor(hass, create_registrations, webhook_client): "name": "Battery State", "state": 100, "type": "sensor", + "entity_category": "diagnostic", "unique_id": "battery_state", + "state_class": "total", "unit_of_measurement": PERCENTAGE, }, }, @@ -38,10 +40,16 @@ async def test_sensor(hass, create_registrations, webhook_client): assert entity.attributes["icon"] == "mdi:battery" assert entity.attributes["unit_of_measurement"] == PERCENTAGE assert entity.attributes["foo"] == "bar" + assert entity.attributes["state_class"] == "total" assert entity.domain == "sensor" assert entity.name == "Test 1 Battery State" assert entity.state == "100" + assert ( + er.async_get(hass).async_get("sensor.test_1_battery_state").entity_category + == "diagnostic" + ) + update_resp = await webhook_client.post( webhook_url, json={ diff --git a/tests/components/mobile_app/test_webhook.py b/tests/components/mobile_app/test_webhook.py index d5cb72fa850..8dc2086c495 100644 --- a/tests/components/mobile_app/test_webhook.py +++ b/tests/components/mobile_app/test_webhook.py @@ -473,3 +473,44 @@ async def test_webhook_handle_scan_tag(hass, create_registrations, webhook_clien assert len(events) == 1 assert events[0].data["tag_id"] == "mock-tag-id" assert events[0].data["device_id"] == "mock-device-id" + + +async def test_register_sensor_limits_state_class( + hass, create_registrations, webhook_client +): + """Test that we limit state classes to sensors only.""" + webhook_id = create_registrations[1]["webhook_id"] + webhook_url = f"/api/webhook/{webhook_id}" + + reg_resp = await webhook_client.post( + webhook_url, + json={ + "type": "register_sensor", + "data": { + "name": "Battery State", + "state": 100, + "type": "sensor", + "state_class": "total", + "unique_id": "abcd", + }, + }, + ) + + assert reg_resp.status == 201 + + reg_resp = await webhook_client.post( + webhook_url, + json={ + "type": "register_sensor", + "data": { + "name": "Battery State", + "state": 100, + "type": "binary_sensor", + "state_class": "total", + "unique_id": "efgh", + }, + }, + ) + + # This means it was ignored. + assert reg_resp.status == 200