Adjust logic related to entity platform state (#147882)

* Adjust logic related to entity platform state

* Break up hard to read if-statement

* Add and improve tests
This commit is contained in:
Erik Montnemery 2025-07-02 14:57:53 +02:00 committed by GitHub
parent 7447cf329b
commit 943fb9948b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 251 additions and 24 deletions

View File

@ -215,16 +215,19 @@ class StateInfo(TypedDict):
class EntityPlatformState(Enum):
"""The platform state of an entity."""
# Not Added: Not yet added to a platform, polling updates
# are written to the state machine.
# Not Added: Not yet added to a platform, states are not written to the
# state machine.
NOT_ADDED = auto()
# Added: Added to a platform, polling updates
# are written to the state machine.
# Adding: Preparing for adding to a platform, states are not written to the
# state machine.
ADDING = auto()
# Added: Added to a platform, states are written to the state machine.
ADDED = auto()
# Removed: Removed from a platform, polling updates
# are not written to the state machine.
# Removed: Removed from a platform, states are not written to the
# state machine.
REMOVED = auto()
@ -1122,21 +1125,24 @@ class Entity(
@callback
def _async_write_ha_state(self) -> None:
"""Write the state to the state machine."""
if self._platform_state is EntityPlatformState.REMOVED:
# Polling returned after the entity has already been removed
return
if (entry := self.registry_entry) and entry.disabled_by:
if not self._disabled_reported:
self._disabled_reported = True
_LOGGER.warning(
(
"Entity %s is incorrectly being triggered for updates while it"
" is disabled. This is a bug in the %s integration"
),
self.entity_id,
self.platform.platform_name,
)
# The check for self.platform guards against integrations not using an
# EntityComponent (which has not been allowed since HA Core 2024.1)
if not self.platform:
if self._platform_state is EntityPlatformState.REMOVED:
# Don't write state if the entity is not added to the platform.
return
elif self._platform_state is not EntityPlatformState.ADDED:
if (entry := self.registry_entry) and entry.disabled_by:
if not self._disabled_reported:
self._disabled_reported = True
_LOGGER.warning(
(
"Entity %s is incorrectly being triggered for updates while it"
" is disabled. This is a bug in the %s integration"
),
self.entity_id,
self.platform.platform_name,
)
return
state_calculate_start = timer()
@ -1145,7 +1151,7 @@ class Entity(
)
time_now = timer()
if entry:
if entry := self.registry_entry:
# Make sure capabilities in the entity registry are up to date. Capabilities
# include capability attributes, device class and supported features
supported_features = supported_features or 0
@ -1346,7 +1352,7 @@ class Entity(
self.hass = hass
self.platform = platform
self.parallel_updates = parallel_updates
self._platform_state = EntityPlatformState.ADDED
self._platform_state = EntityPlatformState.ADDING
def _call_on_remove_callbacks(self) -> None:
"""Call callbacks registered by async_on_remove."""
@ -1370,6 +1376,7 @@ class Entity(
"""Finish adding an entity to a platform."""
await self.async_internal_added_to_hass()
await self.async_added_to_hass()
self._platform_state = EntityPlatformState.ADDED
self.async_write_ha_state()
@final

View File

@ -32,7 +32,7 @@ from homeassistant.core import (
ReleaseChannel,
callback,
)
from homeassistant.exceptions import HomeAssistantError
from homeassistant.exceptions import HomeAssistantError, NoEntitySpecifiedError
from homeassistant.helpers import device_registry as dr, entity, entity_registry as er
from homeassistant.helpers.entity_component import async_update_entity
from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
@ -584,10 +584,13 @@ async def test_async_remove_no_platform(hass: HomeAssistant) -> None:
ent = entity.Entity()
ent.hass = hass
ent.entity_id = "test.test"
assert ent._platform_state == entity.EntityPlatformState.NOT_ADDED
ent.async_write_ha_state()
assert ent._platform_state == entity.EntityPlatformState.NOT_ADDED
assert len(hass.states.async_entity_ids()) == 1
await ent.async_remove()
assert len(hass.states.async_entity_ids()) == 0
assert ent._platform_state == entity.EntityPlatformState.REMOVED
async def test_async_remove_runs_callbacks(hass: HomeAssistant) -> None:
@ -597,10 +600,13 @@ async def test_async_remove_runs_callbacks(hass: HomeAssistant) -> None:
platform = MockEntityPlatform(hass, domain="test")
ent = entity.Entity()
ent.entity_id = "test.test"
assert ent._platform_state == entity.EntityPlatformState.NOT_ADDED
await platform.async_add_entities([ent])
assert ent._platform_state == entity.EntityPlatformState.ADDED
ent.async_on_remove(lambda: result.append(1))
await ent.async_remove()
assert len(result) == 1
assert ent._platform_state == entity.EntityPlatformState.REMOVED
async def test_async_remove_ignores_in_flight_polling(hass: HomeAssistant) -> None:
@ -647,10 +653,12 @@ async def test_async_remove_twice(hass: HomeAssistant) -> None:
await ent.async_remove()
assert len(result) == 1
assert len(ent.remove_calls) == 1
assert ent._platform_state == entity.EntityPlatformState.REMOVED
await ent.async_remove()
assert len(result) == 1
assert len(ent.remove_calls) == 1
assert ent._platform_state == entity.EntityPlatformState.REMOVED
async def test_set_context(hass: HomeAssistant) -> None:
@ -774,6 +782,7 @@ async def test_warn_slow_write_state(
mock_entity.hass = hass
mock_entity.entity_id = "comp_test.test_entity"
mock_entity.platform = MagicMock(platform_name="hue")
mock_entity._platform_state = entity.EntityPlatformState.ADDED
with patch("homeassistant.helpers.entity.timer", side_effect=[0, 10]):
mock_entity.async_write_ha_state()
@ -801,6 +810,7 @@ async def test_warn_slow_write_state_custom_component(
mock_entity.hass = hass
mock_entity.entity_id = "comp_test.test_entity"
mock_entity.platform = MagicMock(platform_name="hue")
mock_entity._platform_state = entity.EntityPlatformState.ADDED
with patch("homeassistant.helpers.entity.timer", side_effect=[0, 10]):
mock_entity.async_write_ha_state()
@ -1781,9 +1791,12 @@ async def test_reuse_entity_object_after_abort(
platform = MockEntityPlatform(hass, domain="test")
ent = entity.Entity()
ent.entity_id = "invalid"
assert ent._platform_state == entity.EntityPlatformState.NOT_ADDED
await platform.async_add_entities([ent])
assert ent._platform_state == entity.EntityPlatformState.REMOVED
assert "Invalid entity ID: invalid" in caplog.text
await platform.async_add_entities([ent])
assert ent._platform_state == entity.EntityPlatformState.REMOVED
assert (
"Entity 'invalid' cannot be added a second time to an entity platform"
in caplog.text
@ -1800,17 +1813,21 @@ async def test_reuse_entity_object_after_entity_registry_remove(
platform = MockEntityPlatform(hass, domain="test", platform_name="test")
ent = entity.Entity()
ent._attr_unique_id = "5678"
assert ent._platform_state == entity.EntityPlatformState.NOT_ADDED
await platform.async_add_entities([ent])
assert ent.registry_entry is entry
assert len(hass.states.async_entity_ids()) == 1
assert ent._platform_state == entity.EntityPlatformState.ADDED
entity_registry.async_remove(entry.entity_id)
await hass.async_block_till_done()
assert len(hass.states.async_entity_ids()) == 0
assert ent._platform_state == entity.EntityPlatformState.REMOVED
await platform.async_add_entities([ent])
assert "Entity 'test.test_5678' cannot be added a second time" in caplog.text
assert len(hass.states.async_entity_ids()) == 0
assert ent._platform_state == entity.EntityPlatformState.REMOVED
async def test_reuse_entity_object_after_entity_registry_disabled(
@ -1823,19 +1840,23 @@ async def test_reuse_entity_object_after_entity_registry_disabled(
platform = MockEntityPlatform(hass, domain="test", platform_name="test")
ent = entity.Entity()
ent._attr_unique_id = "5678"
assert ent._platform_state == entity.EntityPlatformState.NOT_ADDED
await platform.async_add_entities([ent])
assert ent.registry_entry is entry
assert len(hass.states.async_entity_ids()) == 1
assert ent._platform_state == entity.EntityPlatformState.ADDED
entity_registry.async_update_entity(
entry.entity_id, disabled_by=er.RegistryEntryDisabler.USER
)
await hass.async_block_till_done()
assert len(hass.states.async_entity_ids()) == 0
assert ent._platform_state == entity.EntityPlatformState.REMOVED
await platform.async_add_entities([ent])
assert len(hass.states.async_entity_ids()) == 0
assert "Entity 'test.test_5678' cannot be added a second time" in caplog.text
assert ent._platform_state == entity.EntityPlatformState.REMOVED
async def test_change_entity_id(
@ -1865,9 +1886,11 @@ async def test_change_entity_id(
platform = MockEntityPlatform(hass, domain="test")
ent = MockEntity()
assert ent._platform_state == entity.EntityPlatformState.NOT_ADDED
await platform.async_add_entities([ent])
assert hass.states.get("test.test").state == STATE_UNKNOWN
assert len(ent.added_calls) == 1
assert ent._platform_state == entity.EntityPlatformState.ADDED
entry = entity_registry.async_update_entity(
entry.entity_id, new_entity_id="test.test2"
@ -1877,6 +1900,7 @@ async def test_change_entity_id(
assert len(result) == 1
assert len(ent.added_calls) == 2
assert len(ent.remove_calls) == 1
assert ent._platform_state == entity.EntityPlatformState.ADDED
entity_registry.async_update_entity(entry.entity_id, new_entity_id="test.test3")
await hass.async_block_till_done()
@ -1884,6 +1908,7 @@ async def test_change_entity_id(
assert len(result) == 2
assert len(ent.added_calls) == 3
assert len(ent.remove_calls) == 2
assert ent._platform_state == entity.EntityPlatformState.ADDED
def test_entity_description_as_dataclass(snapshot: SnapshotAssertion) -> None:
@ -2524,6 +2549,7 @@ async def test_remove_entity_registry(
assert len(result) == 1
assert len(ent.added_calls) == 1
assert len(ent.remove_calls) == 1
assert ent._platform_state == entity.EntityPlatformState.REMOVED
assert hass.states.get("test.test") is None
@ -2628,6 +2654,7 @@ async def test_async_write_ha_state_thread_safety_always(
ent.entity_id = "test.any"
ent.hass = hass
ent.platform = MockEntityPlatform(hass, domain="test")
ent._platform_state = entity.EntityPlatformState.ADDED
ent.async_write_ha_state()
assert hass.states.get(ent.entity_id)
@ -2641,3 +2668,196 @@ async def test_async_write_ha_state_thread_safety_always(
):
await hass.async_add_executor_job(ent2.async_write_ha_state)
assert not hass.states.get(ent2.entity_id)
async def test_platform_state(
hass: HomeAssistant, entity_registry: er.EntityRegistry
) -> None:
"""Test platform state."""
entry = entity_registry.async_get_or_create(
"test", "test_platform", "5678", suggested_object_id="test"
)
assert entry.entity_id == "test.test"
class MockEntity(entity.Entity):
_attr_unique_id = "5678"
async def async_added_to_hass(self):
# The attempt to write when in state ADDING should be ignored
assert self._platform_state == entity.EntityPlatformState.ADDING
self._attr_state = "added_to_hass"
self.async_write_ha_state()
assert hass.states.get("test.test") is None
async def async_will_remove_from_hass(self):
# The attempt to write when in state REMOVED should be ignored
assert self._platform_state == entity.EntityPlatformState.REMOVED
assert hass.states.get("test.test").state == "added_to_hass"
self._attr_state = "will_remove_from_hass"
self.async_write_ha_state()
assert hass.states.get("test.test").state == "added_to_hass"
platform = MockEntityPlatform(hass, domain="test")
ent = MockEntity()
assert ent._platform_state == entity.EntityPlatformState.NOT_ADDED
await platform.async_add_entities([ent])
assert hass.states.get("test.test").state == "added_to_hass"
assert ent._platform_state == entity.EntityPlatformState.ADDED
entry = entity_registry.async_remove(entry.entity_id)
await hass.async_block_till_done()
assert ent._platform_state == entity.EntityPlatformState.REMOVED
assert hass.states.get("test.test") is None
async def test_platform_state_fail_to_add(
hass: HomeAssistant, entity_registry: er.EntityRegistry
) -> None:
"""Test platform state when raising from async_added_to_hass."""
entry = entity_registry.async_get_or_create(
"test", "test_platform", "5678", suggested_object_id="test"
)
assert entry.entity_id == "test.test"
class MockEntity(entity.Entity):
_attr_unique_id = "5678"
async def async_added_to_hass(self):
raise ValueError("Failed to add entity")
platform = MockEntityPlatform(hass, domain="test")
ent = MockEntity()
assert ent._platform_state == entity.EntityPlatformState.NOT_ADDED
await platform.async_add_entities([ent])
assert hass.states.get("test.test") is None
assert ent._platform_state == entity.EntityPlatformState.ADDING
entry = entity_registry.async_remove(entry.entity_id)
await hass.async_block_till_done()
assert ent._platform_state == entity.EntityPlatformState.REMOVED
assert hass.states.get("test.test") is None
async def test_platform_state_write_from_init(
hass: HomeAssistant, caplog: pytest.LogCaptureFixture
) -> None:
"""Test platform state when an entity attempts to write from init."""
class MockEntity(entity.Entity):
def __init__(self, hass: HomeAssistant) -> None:
self.hass = hass
# The attempt to write when in state NOT_ADDED is prevented because
# the entity has no entity_id set
self._attr_state = "init"
with pytest.raises(NoEntitySpecifiedError):
self.async_write_ha_state()
assert len(hass.states.async_all()) == 0
platform = MockEntityPlatform(hass, domain="test")
ent = MockEntity(hass)
assert ent._platform_state == entity.EntityPlatformState.NOT_ADDED
await platform.async_add_entities([ent])
assert hass.states.get("test.unnamed_device").state == "init"
assert ent._platform_state == entity.EntityPlatformState.ADDED
assert len(hass.states.async_all()) == 1
assert "Platform test_platform does not generate unique IDs." not in caplog.text
assert "Entity id already exists" not in caplog.text
async def test_platform_state_write_from_init_entity_id(
hass: HomeAssistant, caplog: pytest.LogCaptureFixture
) -> None:
"""Test platform state when an entity attempts to write from init.
The outcome of this test is a bit illogical, when we no longer allow
entities without platforms, attempts to write when state is NOT_ADDED
will be blocked.
"""
class MockEntity(entity.Entity):
def __init__(self, hass: HomeAssistant) -> None:
self.entity_id = "test.test"
self.hass = hass
# The attempt to write when in state NOT_ADDED is not prevented because
# the platform is not yet set
assert self._platform_state == entity.EntityPlatformState.NOT_ADDED
self._attr_state = "init"
self.async_write_ha_state()
assert hass.states.get("test.test").state == "init"
async def async_added_to_hass(self):
raise NotImplementedError("Should not be called")
async def async_will_remove_from_hass(self):
raise NotImplementedError("Should not be called")
platform = MockEntityPlatform(hass, domain="test")
ent = MockEntity(hass)
assert ent._platform_state == entity.EntityPlatformState.NOT_ADDED
await platform.async_add_entities([ent])
assert hass.states.get("test.test").state == "init"
assert ent._platform_state == entity.EntityPlatformState.REMOVED
assert len(hass.states.async_all()) == 1
# The early attempt to write is interpreted as a state collision
assert "Platform test_platform does not generate unique IDs." not in caplog.text
assert "Entity id already exists - ignoring: test.test" in caplog.text
async def test_platform_state_write_from_init_unique_id(
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
caplog: pytest.LogCaptureFixture,
) -> None:
"""Test platform state when an entity attempts to write from init.
The outcome of this test is a bit illogical, when we no longer allow
entities without platforms, attempts to write when state is NOT_ADDED
will be blocked.
"""
entry = entity_registry.async_get_or_create(
"test", "test_platform", "5678", suggested_object_id="test"
)
assert entry.entity_id == "test.test"
class MockEntity(entity.Entity):
_attr_unique_id = "5678"
def __init__(self, hass: HomeAssistant) -> None:
self.entity_id = "test.test"
self.hass = hass
# The attempt to write when in state NOT_ADDED is not prevented because
# the platform is not yet set
assert self._platform_state == entity.EntityPlatformState.NOT_ADDED
self._attr_state = "init"
self.async_write_ha_state()
assert hass.states.get("test.test").state == "init"
async def async_added_to_hass(self):
raise NotImplementedError("Should not be called")
async def async_will_remove_from_hass(self):
raise NotImplementedError("Should not be called")
platform = MockEntityPlatform(hass, domain="test")
ent = MockEntity(hass)
assert ent._platform_state == entity.EntityPlatformState.NOT_ADDED
await platform.async_add_entities([ent])
assert hass.states.get("test.test").state == "init"
assert ent._platform_state == entity.EntityPlatformState.REMOVED
assert len(hass.states.async_all()) == 1
# The early attempt to write is interpreted as a unique ID collision
assert "Platform test_platform does not generate unique IDs." in caplog.text
assert "Entity id already exists - ignoring: test.test" not in caplog.text