From a122cbab6135544a9bd42c44b4c3e062f66cf573 Mon Sep 17 00:00:00 2001 From: Paulus Schoutsen Date: Sun, 31 Oct 2021 20:21:46 -0700 Subject: [PATCH] Mobile app to update entity registry on re-register sensors (#58378) Co-authored-by: J. Nick Koston --- .../components/mobile_app/binary_sensor.py | 2 + homeassistant/components/mobile_app/sensor.py | 2 + .../components/mobile_app/webhook.py | 20 +++++ homeassistant/helpers/entity.py | 5 +- homeassistant/helpers/entity_registry.py | 82 +++++++++++-------- tests/components/mobile_app/test_webhook.py | 56 +++++++++++++ 6 files changed, 128 insertions(+), 39 deletions(-) diff --git a/homeassistant/components/mobile_app/binary_sensor.py b/homeassistant/components/mobile_app/binary_sensor.py index 94bcd3b1c0d..2e3681b7618 100644 --- a/homeassistant/components/mobile_app/binary_sensor.py +++ b/homeassistant/components/mobile_app/binary_sensor.py @@ -9,6 +9,7 @@ from .const import ( ATTR_DEVICE_NAME, ATTR_SENSOR_ATTRIBUTES, ATTR_SENSOR_DEVICE_CLASS, + ATTR_SENSOR_ENTITY_CATEGORY, ATTR_SENSOR_ICON, ATTR_SENSOR_NAME, ATTR_SENSOR_STATE, @@ -40,6 +41,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities): ATTR_SENSOR_STATE: None, ATTR_SENSOR_TYPE: entry.domain, ATTR_SENSOR_UNIQUE_ID: entry.unique_id, + ATTR_SENSOR_ENTITY_CATEGORY: entry.entity_category, } entities.append(MobileAppBinarySensor(config, entry.device_id, config_entry)) diff --git a/homeassistant/components/mobile_app/sensor.py b/homeassistant/components/mobile_app/sensor.py index 533e33e84bb..ff4ca491411 100644 --- a/homeassistant/components/mobile_app/sensor.py +++ b/homeassistant/components/mobile_app/sensor.py @@ -11,6 +11,7 @@ from .const import ( ATTR_DEVICE_NAME, ATTR_SENSOR_ATTRIBUTES, ATTR_SENSOR_DEVICE_CLASS, + ATTR_SENSOR_ENTITY_CATEGORY, ATTR_SENSOR_ICON, ATTR_SENSOR_NAME, ATTR_SENSOR_STATE, @@ -45,6 +46,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities): ATTR_SENSOR_TYPE: entry.domain, ATTR_SENSOR_UNIQUE_ID: entry.unique_id, ATTR_SENSOR_UOM: entry.unit_of_measurement, + ATTR_SENSOR_ENTITY_CATEGORY: entry.entity_category, } entities.append(MobileAppSensor(config, entry.device_id, config_entry)) diff --git a/homeassistant/components/mobile_app/webhook.py b/homeassistant/components/mobile_app/webhook.py index 9a0d391373c..ebba383636b 100644 --- a/homeassistant/components/mobile_app/webhook.py +++ b/homeassistant/components/mobile_app/webhook.py @@ -446,6 +446,26 @@ async def webhook_register_sensor(hass, config_entry, data): "Re-register for %s of existing sensor %s", device_name, unique_id ) + entry = entity_registry.async_get(existing_sensor) + changes = {} + + if ( + new_name := f"{device_name} {data[ATTR_SENSOR_NAME]}" + ) != entry.original_name: + changes["original_name"] = new_name + + for ent_reg_key, data_key in ( + ("device_class", ATTR_SENSOR_DEVICE_CLASS), + ("unit_of_measurement", ATTR_SENSOR_UOM), + ("entity_category", ATTR_SENSOR_ENTITY_CATEGORY), + ("original_icon", ATTR_SENSOR_ICON), + ): + if data_key in data and getattr(entry, ent_reg_key) != data[data_key]: + changes[ent_reg_key] = data[data_key] + + if changes: + entity_registry.async_update_entity(existing_sensor, **changes) + async_dispatcher_send(hass, SIGNAL_SENSOR_UPDATE, data) else: register_signal = f"{DOMAIN}_{data[ATTR_SENSOR_TYPE]}_register" diff --git a/homeassistant/helpers/entity.py b/homeassistant/helpers/entity.py index 3b2a12c687e..9ee7221ffa6 100644 --- a/homeassistant/helpers/entity.py +++ b/homeassistant/helpers/entity.py @@ -38,7 +38,6 @@ from homeassistant.core import CALLBACK_TYPE, Context, HomeAssistant, callback from homeassistant.exceptions import HomeAssistantError, NoEntitySpecifiedError from homeassistant.helpers import entity_registry as er from homeassistant.helpers.entity_platform import EntityPlatform -from homeassistant.helpers.entity_registry import RegistryEntry from homeassistant.helpers.event import Event, async_track_entity_registry_updated_event from homeassistant.helpers.typing import StateType from homeassistant.loader import bind_hass @@ -227,7 +226,7 @@ class Entity(ABC): parallel_updates: asyncio.Semaphore | None = None # Entry in the entity registry - registry_entry: RegistryEntry | None = None + registry_entry: er.RegistryEntry | None = None # Hold list for functions to call on remove. _on_remove: list[CALLBACK_TYPE] | None = None @@ -806,7 +805,7 @@ class Entity(ABC): if data["action"] != "update": return - ent_reg = await self.hass.helpers.entity_registry.async_get_registry() + ent_reg = er.async_get(self.hass) old = self.registry_entry self.registry_entry = ent_reg.async_get(data["entity_id"]) assert self.registry_entry is not None diff --git a/homeassistant/helpers/entity_registry.py b/homeassistant/helpers/entity_registry.py index bedbdc51785..b7b0eed2f32 100644 --- a/homeassistant/helpers/entity_registry.py +++ b/homeassistant/helpers/entity_registry.py @@ -243,21 +243,21 @@ class EntityRegistry: unique_id: str, *, # To influence entity ID generation - suggested_object_id: str | None = None, known_object_ids: Iterable[str] | None = None, + suggested_object_id: str | None = None, # To disable an entity if it gets created disabled_by: str | None = None, # Data that we want entry to have - config_entry: ConfigEntry | None = None, - device_id: str | None = None, area_id: str | None = None, capabilities: Mapping[str, Any] | None = None, - supported_features: int | None = None, + config_entry: ConfigEntry | None = None, device_class: str | None = None, - unit_of_measurement: str | None = None, - original_name: str | None = None, - original_icon: str | None = None, + device_id: str | None = None, entity_category: str | None = None, + original_icon: str | None = None, + original_name: str | None = None, + supported_features: int | None = None, + unit_of_measurement: str | None = None, ) -> RegistryEntry: """Get entity. Create if it doesn't exist.""" config_entry_id = None @@ -300,20 +300,20 @@ class EntityRegistry: disabled_by = DISABLED_INTEGRATION entity = RegistryEntry( - entity_id=entity_id, - config_entry_id=config_entry_id, - device_id=device_id, area_id=area_id, - unique_id=unique_id, - platform=platform, - disabled_by=disabled_by, capabilities=capabilities, - supported_features=supported_features or 0, + config_entry_id=config_entry_id, device_class=device_class, - unit_of_measurement=unit_of_measurement, - original_name=original_name, - original_icon=original_icon, + device_id=device_id, + disabled_by=disabled_by, entity_category=entity_category, + entity_id=entity_id, + original_icon=original_icon, + original_name=original_name, + platform=platform, + supported_features=supported_features or 0, + unique_id=unique_id, + unit_of_measurement=unit_of_measurement, ) self._register_entry(entity) _LOGGER.info("Registered new %s.%s entity: %s", domain, platform, entity_id) @@ -383,24 +383,34 @@ class EntityRegistry: self, entity_id: str, *, - name: str | None | UndefinedType = UNDEFINED, - icon: str | None | UndefinedType = UNDEFINED, - config_entry_id: str | None | UndefinedType = UNDEFINED, area_id: str | None | UndefinedType = UNDEFINED, + config_entry_id: str | None | UndefinedType = UNDEFINED, + device_class: str | None | UndefinedType = UNDEFINED, + disabled_by: str | None | UndefinedType = UNDEFINED, + entity_category: str | None | UndefinedType = UNDEFINED, + icon: str | None | UndefinedType = UNDEFINED, + name: str | None | UndefinedType = UNDEFINED, new_entity_id: str | UndefinedType = UNDEFINED, new_unique_id: str | UndefinedType = UNDEFINED, - disabled_by: str | None | UndefinedType = UNDEFINED, + original_icon: str | None | UndefinedType = UNDEFINED, + original_name: str | None | UndefinedType = UNDEFINED, + unit_of_measurement: str | None | UndefinedType = UNDEFINED, ) -> RegistryEntry: """Update properties of an entity.""" return self._async_update_entity( entity_id, - name=name, - icon=icon, - config_entry_id=config_entry_id, area_id=area_id, + config_entry_id=config_entry_id, + device_class=device_class, + disabled_by=disabled_by, + entity_category=entity_category, + icon=icon, + name=name, new_entity_id=new_entity_id, new_unique_id=new_unique_id, - disabled_by=disabled_by, + original_icon=original_icon, + original_name=original_name, + unit_of_measurement=unit_of_measurement, ) @callback @@ -408,21 +418,21 @@ class EntityRegistry: self, entity_id: str, *, - name: str | None | UndefinedType = UNDEFINED, - icon: str | None | UndefinedType = UNDEFINED, - config_entry_id: str | None | UndefinedType = UNDEFINED, - new_entity_id: str | UndefinedType = UNDEFINED, - device_id: str | None | UndefinedType = UNDEFINED, area_id: str | None | UndefinedType = UNDEFINED, - new_unique_id: str | UndefinedType = UNDEFINED, - disabled_by: str | None | UndefinedType = UNDEFINED, capabilities: Mapping[str, Any] | None | UndefinedType = UNDEFINED, - supported_features: int | UndefinedType = UNDEFINED, + config_entry_id: str | None | UndefinedType = UNDEFINED, device_class: str | None | UndefinedType = UNDEFINED, - unit_of_measurement: str | None | UndefinedType = UNDEFINED, - original_name: str | None | UndefinedType = UNDEFINED, - original_icon: str | None | UndefinedType = UNDEFINED, + device_id: str | None | UndefinedType = UNDEFINED, + disabled_by: str | None | UndefinedType = UNDEFINED, entity_category: str | None | UndefinedType = UNDEFINED, + icon: str | None | UndefinedType = UNDEFINED, + name: str | None | UndefinedType = UNDEFINED, + new_entity_id: str | UndefinedType = UNDEFINED, + new_unique_id: str | UndefinedType = UNDEFINED, + original_icon: str | None | UndefinedType = UNDEFINED, + original_name: str | None | UndefinedType = UNDEFINED, + supported_features: int | UndefinedType = UNDEFINED, + unit_of_measurement: str | None | UndefinedType = UNDEFINED, ) -> RegistryEntry: """Private facing update properties method.""" old = self.entities[entity_id] diff --git a/tests/components/mobile_app/test_webhook.py b/tests/components/mobile_app/test_webhook.py index 623abf30e9e..41b939c7113 100644 --- a/tests/components/mobile_app/test_webhook.py +++ b/tests/components/mobile_app/test_webhook.py @@ -10,6 +10,7 @@ from homeassistant.components.zone import DOMAIN as ZONE_DOMAIN from homeassistant.const import CONF_WEBHOOK_ID from homeassistant.core import callback from homeassistant.exceptions import HomeAssistantError +from homeassistant.helpers import entity_registry as er from .const import CALL_SERVICE, FIRE_EVENT, REGISTER_CLEARTEXT, RENDER_TEMPLATE, UPDATE @@ -515,3 +516,58 @@ async def test_register_sensor_limits_state_class( # This means it was ignored. assert reg_resp.status == HTTPStatus.OK + + +async def test_reregister_sensor(hass, create_registrations, webhook_client): + """Test that we can add more info in re-registration.""" + 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", + "unique_id": "abcd", + }, + }, + ) + + assert reg_resp.status == HTTPStatus.CREATED + + ent_reg = er.async_get(hass) + entry = ent_reg.async_get("sensor.test_1_battery_state") + assert entry.original_name == "Test 1 Battery State" + assert entry.device_class is None + assert entry.unit_of_measurement is None + assert entry.entity_category is None + assert entry.original_icon == "mdi:cellphone" + + reg_resp = await webhook_client.post( + webhook_url, + json={ + "type": "register_sensor", + "data": { + "name": "New Name", + "state": 100, + "type": "sensor", + "unique_id": "abcd", + "state_class": "total", + "device_class": "battery", + "entity_category": "diagnostic", + "icon": "mdi:new-icon", + "unit_of_measurement": "%", + }, + }, + ) + + assert reg_resp.status == HTTPStatus.CREATED + entry = ent_reg.async_get("sensor.test_1_battery_state") + assert entry.original_name == "Test 1 New Name" + assert entry.device_class == "battery" + assert entry.unit_of_measurement == "%" + assert entry.entity_category == "diagnostic" + assert entry.original_icon == "mdi:new-icon"