From 943fb9948bdfbc0e38488c4d28bf37bd59d9e0d1 Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Wed, 2 Jul 2025 14:57:53 +0200 Subject: [PATCH] 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 --- homeassistant/helpers/entity.py | 53 ++++---- tests/helpers/test_entity.py | 222 +++++++++++++++++++++++++++++++- 2 files changed, 251 insertions(+), 24 deletions(-) diff --git a/homeassistant/helpers/entity.py b/homeassistant/helpers/entity.py index 39629d07494..352a77af837 100644 --- a/homeassistant/helpers/entity.py +++ b/homeassistant/helpers/entity.py @@ -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 diff --git a/tests/helpers/test_entity.py b/tests/helpers/test_entity.py index 706f1a1a806..24205870779 100644 --- a/tests/helpers/test_entity.py +++ b/tests/helpers/test_entity.py @@ -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