Keep capabilities up to date in the entity registry (#101748)

* Keep capabilities up to date in the entity registry

* Warn if entities update their capabilities very often

* Fix updating of device class

* Stop tracking capability updates once flooding is logged

* Only sync registry if state changed

* Improve test

* Revert "Only sync registry if state changed"

This reverts commit 1c52571596c06444df234d4b088242b494b630f2.

* Avoid calculating device class twice

* Address review comments

* Revert using dataclass

* Fix unintended revert

* Add helper method

tests/helpers/test_entity.py

@@ -8,6 +8,7 @@ import threading
 from typing import Any
 from unittest.mock import MagicMock, PropertyMock, patch
 
+from freezegun.api import FrozenDateTimeFactory
 import pytest
 from syrupy.assertion import SnapshotAssertion
 import voluptuous as vol
@@ -1412,8 +1413,8 @@ async def test_repr_using_stringify_state() -> None:
             """Return the state."""
             raise ValueError("Boom")
 
-    entity = MyEntity(entity_id="test.test", available=False)
-    assert str(entity) == "<entity test.test=unavailable>"
+    my_entity = MyEntity(entity_id="test.test", available=False)
+    assert str(my_entity) == "<entity test.test=unavailable>"
 
 
 async def test_warn_using_async_update_ha_state(
@@ -1761,3 +1762,158 @@ def test_extending_entity_description(snapshot: SnapshotAssertion):
     assert obj == snapshot
     assert obj == CustomInitEntityDescription(key="blah", extra="foo", name="name")
     assert repr(obj) == snapshot
+
+
+async def test_update_capabilities(
+    hass: HomeAssistant,
+    entity_registry: er.EntityRegistry,
+) -> None:
+    """Test entity capabilities are updated automatically."""
+    platform = MockEntityPlatform(hass)
+
+    ent = MockEntity(unique_id="qwer")
+    await platform.async_add_entities([ent])
+
+    entry = entity_registry.async_get(ent.entity_id)
+    assert entry.capabilities is None
+    assert entry.device_class is None
+    assert entry.supported_features == 0
+
+    ent._values["capability_attributes"] = {"bla": "blu"}
+    ent._values["device_class"] = "some_class"
+    ent._values["supported_features"] = 127
+    ent.async_write_ha_state()
+    entry = entity_registry.async_get(ent.entity_id)
+    assert entry.capabilities == {"bla": "blu"}
+    assert entry.original_device_class == "some_class"
+    assert entry.supported_features == 127
+
+    ent._values["capability_attributes"] = None
+    ent._values["device_class"] = None
+    ent._values["supported_features"] = None
+    ent.async_write_ha_state()
+    entry = entity_registry.async_get(ent.entity_id)
+    assert entry.capabilities is None
+    assert entry.original_device_class is None
+    assert entry.supported_features == 0
+
+    # Device class can be overridden by user, make sure that does not break the
+    # automatic updating.
+    entity_registry.async_update_entity(ent.entity_id, device_class="set_by_user")
+    await hass.async_block_till_done()
+    entry = entity_registry.async_get(ent.entity_id)
+    assert entry.capabilities is None
+    assert entry.original_device_class is None
+    assert entry.supported_features == 0
+
+    # This will not trigger a state change because the device class is shadowed
+    # by the entity registry
+    ent._values["device_class"] = "some_class"
+    ent.async_write_ha_state()
+    entry = entity_registry.async_get(ent.entity_id)
+    assert entry.capabilities is None
+    assert entry.original_device_class == "some_class"
+    assert entry.supported_features == 0
+
+
+async def test_update_capabilities_no_unique_id(
+    hass: HomeAssistant,
+    entity_registry: er.EntityRegistry,
+) -> None:
+    """Test entity capabilities are updated automatically."""
+    platform = MockEntityPlatform(hass)
+
+    ent = MockEntity()
+    await platform.async_add_entities([ent])
+
+    assert entity_registry.async_get(ent.entity_id) is None
+
+    ent._values["capability_attributes"] = {"bla": "blu"}
+    ent._values["supported_features"] = 127
+    ent.async_write_ha_state()
+    assert entity_registry.async_get(ent.entity_id) is None
+
+
+async def test_update_capabilities_too_often(
+    hass: HomeAssistant,
+    caplog: pytest.LogCaptureFixture,
+    entity_registry: er.EntityRegistry,
+) -> None:
+    """Test entity capabilities are updated automatically."""
+    capabilities_too_often_warning = "is updating its capabilities too often"
+    platform = MockEntityPlatform(hass)
+
+    ent = MockEntity(unique_id="qwer")
+    await platform.async_add_entities([ent])
+
+    entry = entity_registry.async_get(ent.entity_id)
+    assert entry.capabilities is None
+    assert entry.device_class is None
+    assert entry.supported_features == 0
+
+    for supported_features in range(1, entity.CAPABILITIES_UPDATE_LIMIT + 1):
+        ent._values["capability_attributes"] = {"bla": "blu"}
+        ent._values["device_class"] = "some_class"
+        ent._values["supported_features"] = supported_features
+        ent.async_write_ha_state()
+        entry = entity_registry.async_get(ent.entity_id)
+        assert entry.capabilities == {"bla": "blu"}
+        assert entry.original_device_class == "some_class"
+        assert entry.supported_features == supported_features
+
+    assert capabilities_too_often_warning not in caplog.text
+
+    ent._values["capability_attributes"] = {"bla": "blu"}
+    ent._values["device_class"] = "some_class"
+    ent._values["supported_features"] = supported_features + 1
+    ent.async_write_ha_state()
+    entry = entity_registry.async_get(ent.entity_id)
+    assert entry.capabilities == {"bla": "blu"}
+    assert entry.original_device_class == "some_class"
+    assert entry.supported_features == supported_features + 1
+
+    assert capabilities_too_often_warning in caplog.text
+
+
+async def test_update_capabilities_too_often_cooldown(
+    hass: HomeAssistant,
+    caplog: pytest.LogCaptureFixture,
+    entity_registry: er.EntityRegistry,
+    freezer: FrozenDateTimeFactory,
+) -> None:
+    """Test entity capabilities are updated automatically."""
+    capabilities_too_often_warning = "is updating its capabilities too often"
+    platform = MockEntityPlatform(hass)
+
+    ent = MockEntity(unique_id="qwer")
+    await platform.async_add_entities([ent])
+
+    entry = entity_registry.async_get(ent.entity_id)
+    assert entry.capabilities is None
+    assert entry.device_class is None
+    assert entry.supported_features == 0
+
+    for supported_features in range(1, entity.CAPABILITIES_UPDATE_LIMIT + 1):
+        ent._values["capability_attributes"] = {"bla": "blu"}
+        ent._values["device_class"] = "some_class"
+        ent._values["supported_features"] = supported_features
+        ent.async_write_ha_state()
+        entry = entity_registry.async_get(ent.entity_id)
+        assert entry.capabilities == {"bla": "blu"}
+        assert entry.original_device_class == "some_class"
+        assert entry.supported_features == supported_features
+
+    assert capabilities_too_often_warning not in caplog.text
+
+    freezer.tick(timedelta(minutes=60) + timedelta(seconds=1))
+
+    ent._values["capability_attributes"] = {"bla": "blu"}
+    ent._values["device_class"] = "some_class"
+    ent._values["supported_features"] = supported_features + 1
+    ent.async_write_ha_state()
+    entry = entity_registry.async_get(ent.entity_id)
+    assert entry.capabilities == {"bla": "blu"}
+    assert entry.original_device_class == "some_class"
+    assert entry.supported_features == supported_features + 1
+
+    assert capabilities_too_often_warning not in caplog.text