Prevent polling from recreating an entity after removal (#67750)

tests/helpers/test_entity.py

@@ -545,6 +545,22 @@ async def test_async_remove_runs_callbacks(hass):
     assert len(result) == 1
 
 
+async def test_async_remove_ignores_in_flight_polling(hass):
+    """Test in flight polling is ignored after removing."""
+    result = []
+
+    ent = entity.Entity()
+    ent.hass = hass
+    ent.entity_id = "test.test"
+    ent.async_on_remove(lambda: result.append(1))
+    ent.async_write_ha_state()
+    assert hass.states.get("test.test").state == STATE_UNKNOWN
+    await ent.async_remove()
+    assert len(result) == 1
+    assert hass.states.get("test.test") is None
+    ent.async_write_ha_state()
+
+
 async def test_set_context(hass):
     """Test setting context."""
     context = Context()

tests/helpers/test_entity_platform.py

@@ -390,6 +390,30 @@ async def test_async_remove_with_platform(hass):
     assert len(hass.states.async_entity_ids()) == 0
 
 
+async def test_async_remove_with_platform_update_finishes(hass):
+    """Remove an entity when an update finishes after its been removed."""
+    component = EntityComponent(_LOGGER, DOMAIN, hass)
+    entity1 = MockEntity(name="test_1")
+
+    async def _delayed_update(*args, **kwargs):
+        await asyncio.sleep(0.01)
+
+    entity1.async_update = _delayed_update
+
+    # Add, remove, add, remove and make sure no updates
+    # cause the entity to reappear after removal
+    for i in range(2):
+        await component.async_add_entities([entity1])
+        assert len(hass.states.async_entity_ids()) == 1
+        entity1.async_write_ha_state()
+        assert hass.states.get(entity1.entity_id) is not None
+        task = asyncio.create_task(entity1.async_update_ha_state(True))
+        await entity1.async_remove()
+        assert len(hass.states.async_entity_ids()) == 0
+        await task
+        assert len(hass.states.async_entity_ids()) == 0
+
+
 async def test_not_adding_duplicate_entities_with_unique_id(hass, caplog):
     """Test for not adding duplicate entities."""
     caplog.set_level(logging.ERROR)