Mark entities as unavailable when they are removed but are still registered (#45528)

* Mark entities as unavailable when they are removed but are still registered

* Add sync_entity_lifecycle to collection helper

* Remove debug print

* Lint

* Fix tests

* Fix tests

* Update zha

* Update zone

* Fix tests

* Update hyperion

* Update rfxtrx

* Fix tests

* Pass force_remove=True from integrations

Co-authored-by: Erik <erik@montnemery.com>

homeassistant/components/acmeda/base.py

@@ -32,7 +32,7 @@ class AcmedaBase(entity.Entity):
                 device.id, remove_config_entry_id=self.registry_entry.config_entry_id
             )
 
-        await self.async_remove()
+        await self.async_remove(force_remove=True)
 
     async def async_added_to_hass(self):
         """Entity has been added to hass."""

homeassistant/components/counter/__init__.py

@@ -108,8 +108,8 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool:
     yaml_collection = collection.YamlCollection(
         logging.getLogger(f"{__name__}.yaml_collection"), id_manager
     )
-    collection.attach_entity_component_collection(
+    collection.sync_entity_lifecycle(
-        component, yaml_collection, Counter.from_yaml
+        hass, DOMAIN, DOMAIN, component, yaml_collection, Counter.from_yaml
     )
 
     storage_collection = CounterStorageCollection(
@@ -117,8 +117,8 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool:
         logging.getLogger(f"{__name__}.storage_collection"),
         id_manager,
     )
-    collection.attach_entity_component_collection(
+    collection.sync_entity_lifecycle(
-        component, storage_collection, Counter
+        hass, DOMAIN, DOMAIN, component, storage_collection, Counter
     )
 
     await yaml_collection.async_load(
@@ -130,9 +130,6 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool:
         storage_collection, DOMAIN, DOMAIN, CREATE_FIELDS, UPDATE_FIELDS
     ).async_setup(hass)
 
-    collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, yaml_collection)
-    collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, storage_collection)
-
     component.async_register_entity_service(SERVICE_INCREMENT, {}, "async_increment")
     component.async_register_entity_service(SERVICE_DECREMENT, {}, "async_decrement")
     component.async_register_entity_service(SERVICE_RESET, {}, "async_reset")

homeassistant/components/esphome/__init__.py

@@ -1,5 +1,6 @@
 """Support for esphome devices."""
 import asyncio
+import functools
 import logging
 import math
 from typing import Any, Callable, Dict, List, Optional
@@ -520,7 +521,7 @@ class EsphomeBaseEntity(Entity):
                     f"esphome_{self._entry_id}_remove_"
                     f"{self._component_key}_{self._key}"
                 ),
-                self.async_remove,
+                functools.partial(self.async_remove, force_remove=True),
             )
         )
 

homeassistant/components/gdacs/geo_location.py

@@ -116,7 +116,7 @@ class GdacsEvent(GeolocationEvent):
     @callback
     def _delete_callback(self):
         """Remove this entity."""
-        self.hass.async_create_task(self.async_remove())
+        self.hass.async_create_task(self.async_remove(force_remove=True))
 
     @callback
     def _update_callback(self):

homeassistant/components/geo_json_events/geo_location.py

@@ -144,7 +144,7 @@ class GeoJsonLocationEvent(GeolocationEvent):
         """Remove this entity."""
         self._remove_signal_delete()
         self._remove_signal_update()
-        self.hass.async_create_task(self.async_remove())
+        self.hass.async_create_task(self.async_remove(force_remove=True))
 
     @callback
     def _update_callback(self):

homeassistant/components/geonetnz_quakes/geo_location.py

@@ -102,7 +102,7 @@ class GeonetnzQuakesEvent(GeolocationEvent):
     @callback
     def _delete_callback(self):
         """Remove this entity."""
-        self.hass.async_create_task(self.async_remove())
+        self.hass.async_create_task(self.async_remove(force_remove=True))
 
     @callback
     def _update_callback(self):

homeassistant/components/homematicip_cloud/generic_entity.py

@@ -172,7 +172,7 @@ class HomematicipGenericEntity(Entity):
         """Handle hmip device removal."""
         # Set marker showing that the HmIP device hase been removed.
         self.hmip_device_removed = True
-        self.hass.async_create_task(self.async_remove())
+        self.hass.async_create_task(self.async_remove(force_remove=True))
 
     @property
     def name(self) -> str:

homeassistant/components/hue/helpers.py

@@ -17,7 +17,7 @@ async def remove_devices(bridge, api_ids, current):
         # Device is removed from Hue, so we remove it from Home Assistant
         entity = current[item_id]
         removed_items.append(item_id)
-        await entity.async_remove()
+        await entity.async_remove(force_remove=True)
         ent_registry = await get_ent_reg(bridge.hass)
         if entity.entity_id in ent_registry.entities:
             ent_registry.async_remove(entity.entity_id)

homeassistant/components/hyperion/light.py

@@ -1,6 +1,7 @@
 """Support for Hyperion-NG remotes."""
 from __future__ import annotations
 
+import functools
 import logging
 from types import MappingProxyType
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
@@ -401,7 +402,7 @@ class HyperionBaseLight(LightEntity):
             async_dispatcher_connect(
                 self.hass,
                 SIGNAL_ENTITY_REMOVE.format(self._unique_id),
-                self.async_remove,
+                functools.partial(self.async_remove, force_remove=True),
             )
         )
 

homeassistant/components/hyperion/switch.py

@@ -1,5 +1,6 @@
 """Switch platform for Hyperion."""
 
+import functools
 from typing import Any, Callable, Dict, Optional
 
 from hyperion import client
@@ -199,7 +200,7 @@ class HyperionComponentSwitch(SwitchEntity):
             async_dispatcher_connect(
                 self.hass,
                 SIGNAL_ENTITY_REMOVE.format(self._unique_id),
-                self.async_remove,
+                functools.partial(self.async_remove, force_remove=True),
             )
         )
 

homeassistant/components/ign_sismologia/geo_location.py

@@ -165,7 +165,7 @@ class IgnSismologiaLocationEvent(GeolocationEvent):
         """Remove this entity."""
         self._remove_signal_delete()
         self._remove_signal_update()
-        self.hass.async_create_task(self.async_remove())
+        self.hass.async_create_task(self.async_remove(force_remove=True))
 
     @callback
     def _update_callback(self):

homeassistant/components/input_boolean/__init__.py

@@ -89,8 +89,8 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool:
     yaml_collection = collection.YamlCollection(
         logging.getLogger(f"{__name__}.yaml_collection"), id_manager
     )
-    collection.attach_entity_component_collection(
+    collection.sync_entity_lifecycle(
-        component, yaml_collection, lambda conf: InputBoolean(conf, from_yaml=True)
+        hass, DOMAIN, DOMAIN, component, yaml_collection, InputBoolean.from_yaml
     )
 
     storage_collection = InputBooleanStorageCollection(
@@ -98,8 +98,8 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool:
         logging.getLogger(f"{__name__}.storage_collection"),
         id_manager,
     )
-    collection.attach_entity_component_collection(
+    collection.sync_entity_lifecycle(
-        component, storage_collection, InputBoolean
+        hass, DOMAIN, DOMAIN, component, storage_collection, InputBoolean
     )
 
     await yaml_collection.async_load(
@@ -111,9 +111,6 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool:
         storage_collection, DOMAIN, DOMAIN, CREATE_FIELDS, UPDATE_FIELDS
     ).async_setup(hass)
 
-    collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, yaml_collection)
-    collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, storage_collection)
-
     async def reload_service_handler(service_call: ServiceCallType) -> None:
         """Remove all input booleans and load new ones from config."""
         conf = await component.async_prepare_reload(skip_reset=True)
@@ -146,14 +143,19 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool:
 class InputBoolean(ToggleEntity, RestoreEntity):
     """Representation of a boolean input."""
 
-    def __init__(self, config: typing.Optional[dict], from_yaml: bool = False):
+    def __init__(self, config: typing.Optional[dict]):
         """Initialize a boolean input."""
         self._config = config
-        self._editable = True
+        self.editable = True
         self._state = config.get(CONF_INITIAL)
-        if from_yaml:
+
-            self._editable = False
+    @classmethod
-            self.entity_id = f"{DOMAIN}.{self.unique_id}"
+    def from_yaml(cls, config: typing.Dict) -> "InputBoolean":
+        """Return entity instance initialized from yaml storage."""
+        input_bool = cls(config)
+        input_bool.entity_id = f"{DOMAIN}.{config[CONF_ID]}"
+        input_bool.editable = False
+        return input_bool
 
     @property
     def should_poll(self):
@@ -168,7 +170,7 @@ class InputBoolean(ToggleEntity, RestoreEntity):
     @property
     def state_attributes(self):
         """Return the state attributes of the entity."""
-        return {ATTR_EDITABLE: self._editable}
+        return {ATTR_EDITABLE: self.editable}
 
     @property
     def icon(self):

homeassistant/components/input_datetime/__init__.py

@@ -108,8 +108,8 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool:
     yaml_collection = collection.YamlCollection(
         logging.getLogger(f"{__name__}.yaml_collection"), id_manager
     )
-    collection.attach_entity_component_collection(
+    collection.sync_entity_lifecycle(
-        component, yaml_collection, InputDatetime.from_yaml
+        hass, DOMAIN, DOMAIN, component, yaml_collection, InputDatetime.from_yaml
     )
 
     storage_collection = DateTimeStorageCollection(
@@ -117,8 +117,8 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool:
         logging.getLogger(f"{__name__}.storage_collection"),
         id_manager,
     )
-    collection.attach_entity_component_collection(
+    collection.sync_entity_lifecycle(
-        component, storage_collection, InputDatetime
+        hass, DOMAIN, DOMAIN, component, storage_collection, InputDatetime
     )
 
     await yaml_collection.async_load(
@@ -130,9 +130,6 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool:
         storage_collection, DOMAIN, DOMAIN, CREATE_FIELDS, UPDATE_FIELDS
     ).async_setup(hass)
 
-    collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, yaml_collection)
-    collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, storage_collection)
-
     async def reload_service_handler(service_call: ServiceCallType) -> None:
         """Reload yaml entities."""
         conf = await component.async_prepare_reload(skip_reset=True)

homeassistant/components/input_number/__init__.py

@@ -119,8 +119,8 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool:
     yaml_collection = collection.YamlCollection(
         logging.getLogger(f"{__name__}.yaml_collection"), id_manager
     )
-    collection.attach_entity_component_collection(
+    collection.sync_entity_lifecycle(
-        component, yaml_collection, InputNumber.from_yaml
+        hass, DOMAIN, DOMAIN, component, yaml_collection, InputNumber.from_yaml
     )
 
     storage_collection = NumberStorageCollection(
@@ -128,8 +128,8 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool:
         logging.getLogger(f"{__name__}.storage_collection"),
         id_manager,
     )
-    collection.attach_entity_component_collection(
+    collection.sync_entity_lifecycle(
-        component, storage_collection, InputNumber
+        hass, DOMAIN, DOMAIN, component, storage_collection, InputNumber
     )
 
     await yaml_collection.async_load(
@@ -141,9 +141,6 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool:
         storage_collection, DOMAIN, DOMAIN, CREATE_FIELDS, UPDATE_FIELDS
     ).async_setup(hass)
 
-    collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, yaml_collection)
-    collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, storage_collection)
-
     async def reload_service_handler(service_call: ServiceCallType) -> None:
         """Reload yaml entities."""
         conf = await component.async_prepare_reload(skip_reset=True)

homeassistant/components/input_select/__init__.py

@@ -94,8 +94,8 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool:
     yaml_collection = collection.YamlCollection(
         logging.getLogger(f"{__name__}.yaml_collection"), id_manager
     )
-    collection.attach_entity_component_collection(
+    collection.sync_entity_lifecycle(
-        component, yaml_collection, InputSelect.from_yaml
+        hass, DOMAIN, DOMAIN, component, yaml_collection, InputSelect.from_yaml
     )
 
     storage_collection = InputSelectStorageCollection(
@@ -103,8 +103,8 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool:
         logging.getLogger(f"{__name__}.storage_collection"),
         id_manager,
     )
-    collection.attach_entity_component_collection(
+    collection.sync_entity_lifecycle(
-        component, storage_collection, InputSelect
+        hass, DOMAIN, DOMAIN, component, storage_collection, InputSelect
     )
 
     await yaml_collection.async_load(
@@ -116,9 +116,6 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool:
         storage_collection, DOMAIN, DOMAIN, CREATE_FIELDS, UPDATE_FIELDS
     ).async_setup(hass)
 
-    collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, yaml_collection)
-    collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, storage_collection)
-
     async def reload_service_handler(service_call: ServiceCallType) -> None:
         """Reload yaml entities."""
         conf = await component.async_prepare_reload(skip_reset=True)

homeassistant/components/input_text/__init__.py

@@ -119,8 +119,8 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool:
     yaml_collection = collection.YamlCollection(
         logging.getLogger(f"{__name__}.yaml_collection"), id_manager
     )
-    collection.attach_entity_component_collection(
+    collection.sync_entity_lifecycle(
-        component, yaml_collection, InputText.from_yaml
+        hass, DOMAIN, DOMAIN, component, yaml_collection, InputText.from_yaml
     )
 
     storage_collection = InputTextStorageCollection(
@@ -128,8 +128,8 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool:
         logging.getLogger(f"{__name__}.storage_collection"),
         id_manager,
     )
-    collection.attach_entity_component_collection(
+    collection.sync_entity_lifecycle(
-        component, storage_collection, InputText
+        hass, DOMAIN, DOMAIN, component, storage_collection, InputText
     )
 
     await yaml_collection.async_load(
@@ -141,9 +141,6 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool:
         storage_collection, DOMAIN, DOMAIN, CREATE_FIELDS, UPDATE_FIELDS
     ).async_setup(hass)
 
-    collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, yaml_collection)
-    collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, storage_collection)
-
     async def reload_service_handler(service_call: ServiceCallType) -> None:
         """Reload yaml entities."""
         conf = await component.async_prepare_reload(skip_reset=True)

homeassistant/components/insteon/insteon_entity.py

@@ -1,4 +1,5 @@
 """Insteon base entity."""
+import functools
 import logging
 
 from pyinsteon import devices
@@ -122,7 +123,11 @@ class InsteonEntity(Entity):
         )
         remove_signal = f"{self._insteon_device.address.id}_{SIGNAL_REMOVE_ENTITY}"
         self.async_on_remove(
-            async_dispatcher_connect(self.hass, remove_signal, self.async_remove)
+            async_dispatcher_connect(
+                self.hass,
+                remove_signal,
+                functools.partial(self.async_remove, force_remove=True),
+            )
         )
 
     async def async_will_remove_from_hass(self):

homeassistant/components/mqtt/mixins.py

@@ -387,7 +387,7 @@ class MqttDiscoveryUpdate(Entity):
                 entity_registry.async_remove(self.entity_id)
                 await cleanup_device_registry(self.hass, entity_entry.device_id)
             else:
-                await self.async_remove()
+                await self.async_remove(force_remove=True)
 
         async def discovery_callback(payload):
             """Handle discovery update."""

homeassistant/components/nsw_rural_fire_service_feed/geo_location.py

@@ -210,7 +210,7 @@ class NswRuralFireServiceLocationEvent(GeolocationEvent):
     @callback
     def _delete_callback(self):
         """Remove this entity."""
-        self.hass.async_create_task(self.async_remove())
+        self.hass.async_create_task(self.async_remove(force_remove=True))
 
     @callback
     def _update_callback(self):

homeassistant/components/owntracks/messages.py

@@ -304,7 +304,7 @@ async def async_handle_waypoint(hass, name_base, waypoint):
     if hass.states.get(entity_id) is not None:
         return
 
-    zone = zone_comp.Zone(
+    zone = zone_comp.Zone.from_yaml(
         {
             zone_comp.CONF_NAME: pretty_name,
             zone_comp.CONF_LATITUDE: lat,
@@ -313,7 +313,6 @@ async def async_handle_waypoint(hass, name_base, waypoint):
             zone_comp.CONF_ICON: zone_comp.ICON_IMPORT,
             zone_comp.CONF_PASSIVE: False,
         },
-        False,
     )
     zone.hass = hass
     zone.entity_id = entity_id

homeassistant/components/ozw/entity.py

@@ -268,7 +268,7 @@ class ZWaveDeviceEntity(Entity):
         if not self.values:
             return  # race condition: delete already requested
         if values_id == self.values.values_id:
-            await self.async_remove()
+            await self.async_remove(force_remove=True)
 
 
 def create_device_name(node: OZWNode):

homeassistant/components/person/__init__.py

@@ -306,14 +306,12 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType):
         yaml_collection,
     )
 
-    collection.attach_entity_component_collection(
+    collection.sync_entity_lifecycle(
-        entity_component, yaml_collection, lambda conf: Person(conf, False)
+        hass, DOMAIN, DOMAIN, entity_component, yaml_collection, Person
     )
-    collection.attach_entity_component_collection(
+    collection.sync_entity_lifecycle(
-        entity_component, storage_collection, lambda conf: Person(conf, True)
+        hass, DOMAIN, DOMAIN, entity_component, storage_collection, Person.from_yaml
     )
-    collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, yaml_collection)
-    collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, storage_collection)
 
     await yaml_collection.async_load(
         await filter_yaml_data(hass, config.get(DOMAIN, []))
@@ -358,10 +356,10 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType):
 class Person(RestoreEntity):
     """Represent a tracked person."""
 
-    def __init__(self, config, editable):
+    def __init__(self, config):
         """Set up person."""
         self._config = config
-        self._editable = editable
+        self.editable = True
         self._latitude = None
         self._longitude = None
         self._gps_accuracy = None
@@ -369,6 +367,13 @@ class Person(RestoreEntity):
         self._state = None
         self._unsub_track_device = None
 
+    @classmethod
+    def from_yaml(cls, config):
+        """Return entity instance initialized from yaml storage."""
+        person = cls(config)
+        person.editable = False
+        return person
+
     @property
     def name(self):
         """Return the name of the entity."""
@@ -395,7 +400,7 @@ class Person(RestoreEntity):
     @property
     def state_attributes(self):
         """Return the state attributes of the person."""
-        data = {ATTR_EDITABLE: self._editable, ATTR_ID: self.unique_id}
+        data = {ATTR_EDITABLE: self.editable, ATTR_ID: self.unique_id}
         if self._latitude is not None:
             data[ATTR_LATITUDE] = self._latitude
         if self._longitude is not None:

homeassistant/components/qld_bushfire/geo_location.py

@@ -167,7 +167,7 @@ class QldBushfireLocationEvent(GeolocationEvent):
         """Remove this entity."""
         self._remove_signal_delete()
         self._remove_signal_update()
-        self.hass.async_create_task(self.async_remove())
+        self.hass.async_create_task(self.async_remove(force_remove=True))
 
     @callback
     def _update_callback(self):

homeassistant/components/rfxtrx/__init__.py

@@ -3,6 +3,7 @@ import asyncio
 import binascii
 from collections import OrderedDict
 import copy
+import functools
 import logging
 
 import RFXtrx as rfxtrxmod
@@ -488,7 +489,8 @@ class RfxtrxEntity(RestoreEntity):
 
         self.async_on_remove(
             self.hass.helpers.dispatcher.async_dispatcher_connect(
-                f"{DOMAIN}_{CONF_REMOVE_DEVICE}_{self._device_id}", self.async_remove
+                f"{DOMAIN}_{CONF_REMOVE_DEVICE}_{self._device_id}",
+                functools.partial(self.async_remove, force_remove=True),
             )
         )
 

homeassistant/components/seventeentrack/sensor.py

@@ -244,7 +244,7 @@ class SeventeenTrackPackageSensor(Entity):
 
     async def _remove(self, *_):
         """Remove entity itself."""
-        await self.async_remove()
+        await self.async_remove(force_remove=True)
 
         reg = await self.hass.helpers.entity_registry.async_get_registry()
         entity_id = reg.async_get_entity_id(

homeassistant/components/timer/__init__.py

@@ -107,8 +107,8 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool:
     yaml_collection = collection.YamlCollection(
         logging.getLogger(f"{__name__}.yaml_collection"), id_manager
     )
-    collection.attach_entity_component_collection(
+    collection.sync_entity_lifecycle(
-        component, yaml_collection, Timer.from_yaml
+        hass, DOMAIN, DOMAIN, component, yaml_collection, Timer.from_yaml
     )
 
     storage_collection = TimerStorageCollection(
@@ -116,7 +116,9 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool:
         logging.getLogger(f"{__name__}.storage_collection"),
         id_manager,
     )
-    collection.attach_entity_component_collection(component, storage_collection, Timer)
+    collection.sync_entity_lifecycle(
+        hass, DOMAIN, DOMAIN, component, storage_collection, Timer
+    )
 
     await yaml_collection.async_load(
         [{CONF_ID: id_, **cfg} for id_, cfg in config.get(DOMAIN, {}).items()]
@@ -127,9 +129,6 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool:
         storage_collection, DOMAIN, DOMAIN, CREATE_FIELDS, UPDATE_FIELDS
     ).async_setup(hass)
 
-    collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, yaml_collection)
-    collection.attach_entity_registry_cleaner(hass, DOMAIN, DOMAIN, storage_collection)
-
     async def reload_service_handler(service_call: ServiceCallType) -> None:
         """Reload yaml entities."""
         conf = await component.async_prepare_reload(skip_reset=True)

homeassistant/components/tuya/__init__.py

@@ -392,7 +392,7 @@ class TuyaDevice(Entity):
                 entity_registry.async_remove(self.entity_id)
                 await cleanup_device_registry(self.hass, entity_entry.device_id)
             else:
-                await self.async_remove()
+                await self.async_remove(force_remove=True)
 
     @callback
     def _update_callback(self):

homeassistant/components/unifi/unifi_entity_base.py

@@ -91,7 +91,7 @@ class UniFiBase(Entity):
         entity_registry = await self.hass.helpers.entity_registry.async_get_registry()
         entity_entry = entity_registry.async_get(self.entity_id)
         if not entity_entry:
-            await self.async_remove()
+            await self.async_remove(force_remove=True)
             return
 
         device_registry = await self.hass.helpers.device_registry.async_get_registry()

homeassistant/components/usgs_earthquakes_feed/geo_location.py

@@ -210,7 +210,7 @@ class UsgsEarthquakesEvent(GeolocationEvent):
         """Remove this entity."""
         self._remove_signal_delete()
         self._remove_signal_update()
-        self.hass.async_create_task(self.async_remove())
+        self.hass.async_create_task(self.async_remove(force_remove=True))
 
     @callback
     def _update_callback(self):

homeassistant/components/wled/light.py

@@ -442,7 +442,7 @@ async def async_remove_entity(
 ) -> None:
     """Remove WLED segment light from Home Assistant."""
     entity = current[index]
-    await entity.async_remove()
+    await entity.async_remove(force_remove=True)
     registry = await async_get_entity_registry(coordinator.hass)
     if entity.entity_id in registry.entities:
         registry.async_remove(entity.entity_id)

homeassistant/components/zha/entity.py

@@ -1,6 +1,7 @@
 """Entity for Zigbee Home Automation."""
 
 import asyncio
+import functools
 import logging
 from typing import Any, Awaitable, Dict, List, Optional
 
@@ -165,7 +166,7 @@ class ZhaEntity(BaseZhaEntity, RestoreEntity):
         self.async_accept_signal(
             None,
             f"{SIGNAL_REMOVE}_{self.zha_device.ieee}",
-            self.async_remove,
+            functools.partial(self.async_remove, force_remove=True),
             signal_override=True,
         )
 
@@ -239,7 +240,7 @@ class ZhaGroupEntity(BaseZhaEntity):
             return
 
         self._handled_group_membership = True
-        await self.async_remove()
+        await self.async_remove(force_remove=True)
 
     async def async_added_to_hass(self) -> None:
         """Register callbacks."""

homeassistant/components/zone/__init__.py

@@ -25,7 +25,6 @@ from homeassistant.helpers import (
     config_validation as cv,
     entity,
     entity_component,
-    entity_registry,
     service,
     storage,
 )
@@ -183,8 +182,8 @@ async def async_setup(hass: HomeAssistant, config: Dict) -> bool:
     yaml_collection = collection.IDLessCollection(
         logging.getLogger(f"{__name__}.yaml_collection"), id_manager
     )
-    collection.attach_entity_component_collection(
+    collection.sync_entity_lifecycle(
-        component, yaml_collection, lambda conf: Zone(conf, False)
+        hass, DOMAIN, DOMAIN, component, yaml_collection, Zone.from_yaml
     )
 
     storage_collection = ZoneStorageCollection(
@@ -192,8 +191,8 @@ async def async_setup(hass: HomeAssistant, config: Dict) -> bool:
         logging.getLogger(f"{__name__}.storage_collection"),
         id_manager,
     )
-    collection.attach_entity_component_collection(
+    collection.sync_entity_lifecycle(
-        component, storage_collection, lambda conf: Zone(conf, True)
+        hass, DOMAIN, DOMAIN, component, storage_collection, Zone
     )
 
     if config[DOMAIN]:
@@ -205,18 +204,6 @@ async def async_setup(hass: HomeAssistant, config: Dict) -> bool:
         storage_collection, DOMAIN, DOMAIN, CREATE_FIELDS, UPDATE_FIELDS
     ).async_setup(hass)
 
-    async def _collection_changed(change_type: str, item_id: str, config: Dict) -> None:
-        """Handle a collection change: clean up entity registry on removals."""
-        if change_type != collection.CHANGE_REMOVED:
-            return
-
-        ent_reg = await entity_registry.async_get_registry(hass)
-        ent_reg.async_remove(
-            cast(str, ent_reg.async_get_entity_id(DOMAIN, DOMAIN, item_id))
-        )
-
-    storage_collection.async_add_listener(_collection_changed)
-
     async def reload_service_handler(service_call: ServiceCall) -> None:
         """Remove all zones and load new ones from config."""
         conf = await component.async_prepare_reload(skip_reset=True)
@@ -235,10 +222,7 @@ async def async_setup(hass: HomeAssistant, config: Dict) -> bool:
     if component.get_entity("zone.home"):
         return True
 
-    home_zone = Zone(
+    home_zone = Zone(_home_conf(hass))
-        _home_conf(hass),
-        True,
-    )
     home_zone.entity_id = ENTITY_ID_HOME
     await component.async_add_entities([home_zone])
 
@@ -293,13 +277,21 @@ async def async_unload_entry(
 class Zone(entity.Entity):
     """Representation of a Zone."""
 
-    def __init__(self, config: Dict, editable: bool):
+    def __init__(self, config: Dict):
         """Initialize the zone."""
         self._config = config
-        self._editable = editable
+        self.editable = True
         self._attrs: Optional[Dict] = None
         self._generate_attrs()
 
+    @classmethod
+    def from_yaml(cls, config: Dict) -> "Zone":
+        """Return entity instance initialized from yaml storage."""
+        zone = cls(config)
+        zone.editable = False
+        zone._generate_attrs()  # pylint:disable=protected-access
+        return zone
+
     @property
     def state(self) -> str:
         """Return the state property really does nothing for a zone."""
@@ -346,5 +338,5 @@ class Zone(entity.Entity):
             ATTR_LONGITUDE: self._config[CONF_LONGITUDE],
             ATTR_RADIUS: self._config[CONF_RADIUS],
             ATTR_PASSIVE: self._config[CONF_PASSIVE],
-            ATTR_EDITABLE: self._editable,
+            ATTR_EDITABLE: self.editable,
         }

homeassistant/components/zwave/node_entity.py

@@ -95,7 +95,7 @@ class ZWaveBaseEntity(Entity):
         """Remove this entity and add it back."""
 
         async def _async_remove_and_add():
-            await self.async_remove()
+            await self.async_remove(force_remove=True)
             self.entity_id = None
             await self.platform.async_add_entities([self])
 
@@ -104,7 +104,7 @@ class ZWaveBaseEntity(Entity):
 
     async def node_removed(self):
         """Call when a node is removed from the Z-Wave network."""
-        await self.async_remove()
+        await self.async_remove(force_remove=True)
 
         registry = await async_get_registry(self.hass)
         if self.entity_id not in registry.entities:

homeassistant/helpers/collection.py

@@ -301,7 +301,10 @@ class IDLessCollection(ObservableCollection):
 
 
 @callback
-def attach_entity_component_collection(
+def sync_entity_lifecycle(
+    hass: HomeAssistantType,
+    domain: str,
+    platform: str,
     entity_component: EntityComponent,
     collection: ObservableCollection,
     create_entity: Callable[[dict], Entity],
@@ -318,8 +321,13 @@ def attach_entity_component_collection(
             return
 
         if change_type == CHANGE_REMOVED:
-            entity = entities.pop(item_id)
+            ent_reg = await entity_registry.async_get_registry(hass)
-            await entity.async_remove()
+            ent_to_remove = ent_reg.async_get_entity_id(domain, platform, item_id)
+            if ent_to_remove is not None:
+                ent_reg.async_remove(ent_to_remove)
+            else:
+                await entities[item_id].async_remove(force_remove=True)
+            entities.pop(item_id)
             return
 
         # CHANGE_UPDATED
@@ -328,28 +336,6 @@ def attach_entity_component_collection(
     collection.async_add_listener(_collection_changed)
 
 
-@callback
-def attach_entity_registry_cleaner(
-    hass: HomeAssistantType,
-    domain: str,
-    platform: str,
-    collection: ObservableCollection,
-) -> None:
-    """Attach a listener to clean up entity registry on collection changes."""
-
-    async def _collection_changed(change_type: str, item_id: str, config: Dict) -> None:
-        """Handle a collection change: clean up entity registry on removals."""
-        if change_type != CHANGE_REMOVED:
-            return
-
-        ent_reg = await entity_registry.async_get_registry(hass)
-        ent_to_remove = ent_reg.async_get_entity_id(domain, platform, item_id)
-        if ent_to_remove is not None:
-            ent_reg.async_remove(ent_to_remove)
-
-    collection.async_add_listener(_collection_changed)
-
-
 class StorageCollectionWebsocket:
     """Class to expose storage collection management over websocket."""
 

homeassistant/helpers/entity.py

@@ -530,8 +530,16 @@ class Entity(ABC):
         await self.async_added_to_hass()
         self.async_write_ha_state()
 
-    async def async_remove(self) -> None:
+    async def async_remove(self, *, force_remove: bool = False) -> None:
-        """Remove entity from Home Assistant."""
+        """Remove entity from Home Assistant.
+
+        If the entity has a non disabled entry in the entity registry,
+        the entity's state will be set to unavailable, in the same way
+        as when the entity registry is loaded.
+
+        If the entity doesn't have a non disabled entry in the entity registry,
+        or if force_remove=True, its state will be removed.
+        """
         assert self.hass is not None
 
         if self.platform and not self._added:
@@ -548,7 +556,16 @@ class Entity(ABC):
         await self.async_internal_will_remove_from_hass()
         await self.async_will_remove_from_hass()
 
-        self.hass.states.async_remove(self.entity_id, context=self._context)
+        # Check if entry still exists in entity registry (e.g. unloading config entry)
+        if (
+            not force_remove
+            and self.registry_entry
+            and not self.registry_entry.disabled
+        ):
+            # Set the entity's state will to unavailable + ATTR_RESTORED: True
+            self.registry_entry.write_unavailable_state(self.hass)
+        else:
+            self.hass.states.async_remove(self.entity_id, context=self._context)
 
     async def async_added_to_hass(self) -> None:
         """Run when entity about to be added to hass.
@@ -606,6 +623,7 @@ class Entity(ABC):
         data = event.data
         if data["action"] == "remove":
             await self.async_removed_from_registry()
+            self.registry_entry = None
             await self.async_remove()
 
         if data["action"] != "update":
@@ -617,7 +635,7 @@ class Entity(ABC):
         self.registry_entry = ent_reg.async_get(data["entity_id"])
         assert self.registry_entry is not None
 
-        if self.registry_entry.disabled_by is not None:
+        if self.registry_entry.disabled:
             await self.async_remove()
             return
 
@@ -626,7 +644,7 @@ class Entity(ABC):
             self.async_write_ha_state()
             return
 
-        await self.async_remove()
+        await self.async_remove(force_remove=True)
 
         assert self.platform is not None
         self.entity_id = self.registry_entry.entity_id

homeassistant/helpers/entity_platform.py

@@ -517,7 +517,7 @@ class EntityPlatform:
         if not self.entities:
             return
 
-        tasks = [self.async_remove_entity(entity_id) for entity_id in self.entities]
+        tasks = [entity.async_remove() for entity in self.entities.values()]
 
         await asyncio.gather(*tasks)
 

homeassistant/helpers/entity_registry.py

@@ -115,6 +115,33 @@ class RegistryEntry:
         """Return if entry is disabled."""
         return self.disabled_by is not None
 
+    @callback
+    def write_unavailable_state(self, hass: HomeAssistantType) -> None:
+        """Write the unavailable state to the state machine."""
+        attrs: Dict[str, Any] = {ATTR_RESTORED: True}
+
+        if self.capabilities is not None:
+            attrs.update(self.capabilities)
+
+        if self.supported_features is not None:
+            attrs[ATTR_SUPPORTED_FEATURES] = self.supported_features
+
+        if self.device_class is not None:
+            attrs[ATTR_DEVICE_CLASS] = self.device_class
+
+        if self.unit_of_measurement is not None:
+            attrs[ATTR_UNIT_OF_MEASUREMENT] = self.unit_of_measurement
+
+        name = self.name or self.original_name
+        if name is not None:
+            attrs[ATTR_FRIENDLY_NAME] = name
+
+        icon = self.icon or self.original_icon
+        if icon is not None:
+            attrs[ATTR_ICON] = icon
+
+        hass.states.async_set(self.entity_id, STATE_UNAVAILABLE, attrs)
+
 
 class EntityRegistry:
     """Class to hold a registry of entities."""
@@ -616,36 +643,13 @@ def async_setup_entity_restore(
     @callback
     def _write_unavailable_states(_: Event) -> None:
         """Make sure state machine contains entry for each registered entity."""
-        states = hass.states
+        existing = set(hass.states.async_entity_ids())
-        existing = set(states.async_entity_ids())
 
         for entry in registry.entities.values():
             if entry.entity_id in existing or entry.disabled:
                 continue
 
-            attrs: Dict[str, Any] = {ATTR_RESTORED: True}
+            entry.write_unavailable_state(hass)
-
-            if entry.capabilities is not None:
-                attrs.update(entry.capabilities)
-
-            if entry.supported_features is not None:
-                attrs[ATTR_SUPPORTED_FEATURES] = entry.supported_features
-
-            if entry.device_class is not None:
-                attrs[ATTR_DEVICE_CLASS] = entry.device_class
-
-            if entry.unit_of_measurement is not None:
-                attrs[ATTR_UNIT_OF_MEASUREMENT] = entry.unit_of_measurement
-
-            name = entry.name or entry.original_name
-            if name is not None:
-                attrs[ATTR_FRIENDLY_NAME] = name
-
-            icon = entry.icon or entry.original_icon
-            if icon is not None:
-                attrs[ATTR_ICON] = icon
-
-            states.async_set(entry.entity_id, STATE_UNAVAILABLE, attrs)
 
     hass.bus.async_listen(EVENT_HOMEASSISTANT_START, _write_unavailable_states)
 

tests/components/cert_expiry/test_init.py

@@ -5,7 +5,12 @@ from unittest.mock import patch
 from homeassistant.components.cert_expiry.const import DOMAIN
 from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
 from homeassistant.config_entries import ENTRY_STATE_LOADED, ENTRY_STATE_NOT_LOADED
-from homeassistant.const import CONF_HOST, CONF_PORT, EVENT_HOMEASSISTANT_START
+from homeassistant.const import (
+    CONF_HOST,
+    CONF_PORT,
+    EVENT_HOMEASSISTANT_START,
+    STATE_UNAVAILABLE,
+)
 from homeassistant.setup import async_setup_component
 import homeassistant.util.dt as dt_util
 
@@ -94,4 +99,9 @@ async def test_unload_config_entry(mock_now, hass):
 
     assert entry.state == ENTRY_STATE_NOT_LOADED
     state = hass.states.get("sensor.cert_expiry_timestamp_example_com")
+    assert state.state == STATE_UNAVAILABLE
+
+    await hass.config_entries.async_remove(entry.entry_id)
+    await hass.async_block_till_done()
+    state = hass.states.get("sensor.cert_expiry_timestamp_example_com")
     assert state is None

tests/components/deconz/test_binary_sensor.py

@@ -16,7 +16,7 @@ from homeassistant.components.deconz.const import (
 )
 from homeassistant.components.deconz.gateway import get_gateway_from_config_entry
 from homeassistant.components.deconz.services import SERVICE_DEVICE_REFRESH
-from homeassistant.const import STATE_OFF, STATE_ON
+from homeassistant.const import STATE_OFF, STATE_ON, STATE_UNAVAILABLE
 from homeassistant.helpers.entity_registry import async_entries_for_config_entry
 from homeassistant.setup import async_setup_component
 
@@ -111,6 +111,10 @@ async def test_binary_sensors(hass):
 
     await hass.config_entries.async_unload(config_entry.entry_id)
 
+    assert hass.states.get("binary_sensor.presence_sensor").state == STATE_UNAVAILABLE
+
+    await hass.config_entries.async_remove(config_entry.entry_id)
+    await hass.async_block_till_done()
     assert len(hass.states.async_all()) == 0
 
 

tests/components/deconz/test_climate.py

@@ -39,7 +39,12 @@ from homeassistant.components.deconz.const import (
     DOMAIN as DECONZ_DOMAIN,
 )
 from homeassistant.components.deconz.gateway import get_gateway_from_config_entry
-from homeassistant.const import ATTR_ENTITY_ID, ATTR_TEMPERATURE, STATE_OFF
+from homeassistant.const import (
+    ATTR_ENTITY_ID,
+    ATTR_TEMPERATURE,
+    STATE_OFF,
+    STATE_UNAVAILABLE,
+)
 from homeassistant.setup import async_setup_component
 
 from .test_gateway import DECONZ_WEB_REQUEST, setup_deconz_integration
@@ -361,6 +366,13 @@ async def test_climate_device_without_cooling_support(hass):
 
     await hass.config_entries.async_unload(config_entry.entry_id)
 
+    states = hass.states.async_all()
+    assert len(hass.states.async_all()) == 2
+    for state in states:
+        assert state.state == STATE_UNAVAILABLE
+
+    await hass.config_entries.async_remove(config_entry.entry_id)
+    await hass.async_block_till_done()
     assert len(hass.states.async_all()) == 0
 
 

tests/components/deconz/test_cover.py

@@ -19,7 +19,12 @@ from homeassistant.components.cover import (
 )
 from homeassistant.components.deconz.const import DOMAIN as DECONZ_DOMAIN
 from homeassistant.components.deconz.gateway import get_gateway_from_config_entry
-from homeassistant.const import ATTR_ENTITY_ID, STATE_CLOSED, STATE_OPEN
+from homeassistant.const import (
+    ATTR_ENTITY_ID,
+    STATE_CLOSED,
+    STATE_OPEN,
+    STATE_UNAVAILABLE,
+)
 from homeassistant.setup import async_setup_component
 
 from .test_gateway import DECONZ_WEB_REQUEST, setup_deconz_integration
@@ -251,6 +256,13 @@ async def test_cover(hass):
 
     await hass.config_entries.async_unload(config_entry.entry_id)
 
+    states = hass.states.async_all()
+    assert len(hass.states.async_all()) == 5
+    for state in states:
+        assert state.state == STATE_UNAVAILABLE
+
+    await hass.config_entries.async_remove(config_entry.entry_id)
+    await hass.async_block_till_done()
     assert len(hass.states.async_all()) == 0
 
 

tests/components/deconz/test_deconz_event.py

@@ -4,6 +4,7 @@ from copy import deepcopy
 
 from homeassistant.components.deconz.deconz_event import CONF_DECONZ_EVENT
 from homeassistant.components.deconz.gateway import get_gateway_from_config_entry
+from homeassistant.const import STATE_UNAVAILABLE
 
 from .test_gateway import DECONZ_WEB_REQUEST, setup_deconz_integration
 
@@ -121,5 +122,13 @@ async def test_deconz_events(hass):
 
     await hass.config_entries.async_unload(config_entry.entry_id)
 
+    states = hass.states.async_all()
+    assert len(hass.states.async_all()) == 3
+    for state in states:
+        assert state.state == STATE_UNAVAILABLE
+    assert len(gateway.events) == 0
+
+    await hass.config_entries.async_remove(config_entry.entry_id)
+    await hass.async_block_till_done()
     assert len(hass.states.async_all()) == 0
     assert len(gateway.events) == 0

tests/components/deconz/test_fan.py

@@ -18,7 +18,7 @@ from homeassistant.components.fan import (
     SPEED_MEDIUM,
     SPEED_OFF,
 )
-from homeassistant.const import ATTR_ENTITY_ID, STATE_OFF, STATE_ON
+from homeassistant.const import ATTR_ENTITY_ID, STATE_OFF, STATE_ON, STATE_UNAVAILABLE
 from homeassistant.setup import async_setup_component
 
 from .test_gateway import DECONZ_WEB_REQUEST, setup_deconz_integration
@@ -207,4 +207,11 @@ async def test_fans(hass):
 
     await hass.config_entries.async_unload(config_entry.entry_id)
 
+    states = hass.states.async_all()
+    assert len(hass.states.async_all()) == 2
+    for state in states:
+        assert state.state == STATE_UNAVAILABLE
+
+    await hass.config_entries.async_remove(config_entry.entry_id)
+    await hass.async_block_till_done()
     assert len(hass.states.async_all()) == 0

tests/components/deconz/test_light.py

@@ -31,6 +31,7 @@ from homeassistant.const import (
     ATTR_SUPPORTED_FEATURES,
     STATE_OFF,
     STATE_ON,
+    STATE_UNAVAILABLE,
 )
 from homeassistant.setup import async_setup_component
 
@@ -296,6 +297,13 @@ async def test_lights_and_groups(hass):
 
     await hass.config_entries.async_unload(config_entry.entry_id)
 
+    states = hass.states.async_all()
+    assert len(hass.states.async_all()) == 6
+    for state in states:
+        assert state.state == STATE_UNAVAILABLE
+
+    await hass.config_entries.async_remove(config_entry.entry_id)
+    await hass.async_block_till_done()
     assert len(hass.states.async_all()) == 0
 
 

tests/components/deconz/test_lock.py

@@ -10,7 +10,12 @@ from homeassistant.components.lock import (
     SERVICE_LOCK,
     SERVICE_UNLOCK,
 )
-from homeassistant.const import ATTR_ENTITY_ID, STATE_LOCKED, STATE_UNLOCKED
+from homeassistant.const import (
+    ATTR_ENTITY_ID,
+    STATE_LOCKED,
+    STATE_UNAVAILABLE,
+    STATE_UNLOCKED,
+)
 from homeassistant.setup import async_setup_component
 
 from .test_gateway import DECONZ_WEB_REQUEST, setup_deconz_integration
@@ -104,4 +109,11 @@ async def test_locks(hass):
 
     await hass.config_entries.async_unload(config_entry.entry_id)
 
+    states = hass.states.async_all()
+    assert len(hass.states.async_all()) == 1
+    for state in states:
+        assert state.state == STATE_UNAVAILABLE
+
+    await hass.config_entries.async_remove(config_entry.entry_id)
+    await hass.async_block_till_done()
     assert len(hass.states.async_all()) == 0

tests/components/deconz/test_sensor.py

@@ -12,6 +12,7 @@ from homeassistant.const import (
     DEVICE_CLASS_BATTERY,
     DEVICE_CLASS_ILLUMINANCE,
     DEVICE_CLASS_POWER,
+    STATE_UNAVAILABLE,
 )
 from homeassistant.setup import async_setup_component
 
@@ -165,6 +166,13 @@ async def test_sensors(hass):
 
     await hass.config_entries.async_unload(config_entry.entry_id)
 
+    states = hass.states.async_all()
+    assert len(hass.states.async_all()) == 5
+    for state in states:
+        assert state.state == STATE_UNAVAILABLE
+
+    await hass.config_entries.async_remove(config_entry.entry_id)
+    await hass.async_block_till_done()
     assert len(hass.states.async_all()) == 0
 
 

tests/components/deconz/test_switch.py

@@ -10,7 +10,7 @@ from homeassistant.components.switch import (
     SERVICE_TURN_OFF,
     SERVICE_TURN_ON,
 )
-from homeassistant.const import ATTR_ENTITY_ID, STATE_OFF, STATE_ON
+from homeassistant.const import ATTR_ENTITY_ID, STATE_OFF, STATE_ON, STATE_UNAVAILABLE
 from homeassistant.setup import async_setup_component
 
 from .test_gateway import DECONZ_WEB_REQUEST, setup_deconz_integration
@@ -139,6 +139,13 @@ async def test_power_plugs(hass):
 
     await hass.config_entries.async_unload(config_entry.entry_id)
 
+    states = hass.states.async_all()
+    assert len(hass.states.async_all()) == 4
+    for state in states:
+        assert state.state == STATE_UNAVAILABLE
+
+    await hass.config_entries.async_remove(config_entry.entry_id)
+    await hass.async_block_till_done()
     assert len(hass.states.async_all()) == 0
 
 
@@ -202,4 +209,11 @@ async def test_sirens(hass):
 
     await hass.config_entries.async_unload(config_entry.entry_id)
 
+    states = hass.states.async_all()
+    assert len(hass.states.async_all()) == 2
+    for state in states:
+        assert state.state == STATE_UNAVAILABLE
+
+    await hass.config_entries.async_remove(config_entry.entry_id)
+    await hass.async_block_till_done()
     assert len(hass.states.async_all()) == 0

tests/components/dynalite/test_light.py

@@ -4,7 +4,11 @@ from dynalite_devices_lib.light import DynaliteChannelLightDevice
 import pytest
 
 from homeassistant.components.light import SUPPORT_BRIGHTNESS
-from homeassistant.const import ATTR_FRIENDLY_NAME, ATTR_SUPPORTED_FEATURES
+from homeassistant.const import (
+    ATTR_FRIENDLY_NAME,
+    ATTR_SUPPORTED_FEATURES,
+    STATE_UNAVAILABLE,
+)
 
 from .common import (
     ATTR_METHOD,
@@ -40,11 +44,21 @@ async def test_light_setup(hass, mock_device):
     )
 
 
-async def test_remove_entity(hass, mock_device):
+async def test_unload_config_entry(hass, mock_device):
-    """Test when an entity is removed from HA."""
+    """Test when a config entry is unloaded from HA."""
     await create_entity_from_device(hass, mock_device)
     assert hass.states.get("light.name")
     entry_id = await get_entry_id_from_hass(hass)
     assert await hass.config_entries.async_unload(entry_id)
     await hass.async_block_till_done()
+    assert hass.states.get("light.name").state == STATE_UNAVAILABLE
+
+
+async def test_remove_config_entry(hass, mock_device):
+    """Test when a config entry is removed from HA."""
+    await create_entity_from_device(hass, mock_device)
+    assert hass.states.get("light.name")
+    entry_id = await get_entry_id_from_hass(hass)
+    assert await hass.config_entries.async_remove(entry_id)
+    await hass.async_block_till_done()
     assert not hass.states.get("light.name")

tests/components/eafm/test_sensor.py

@@ -5,7 +5,7 @@ import aiohttp
 import pytest
 
 from homeassistant import config_entries
-from homeassistant.const import ATTR_UNIT_OF_MEASUREMENT
+from homeassistant.const import ATTR_UNIT_OF_MEASUREMENT, STATE_UNAVAILABLE
 from homeassistant.setup import async_setup_component
 import homeassistant.util.dt as dt_util
 
@@ -428,5 +428,8 @@ async def test_unload_entry(hass, mock_get_station):
 
     assert await entry.async_unload(hass)
 
-    # And the entity should be gone
+    # And the entity should be unavailable
-    assert not hass.states.get("sensor.my_station_water_level_stage")
+    assert (
+        hass.states.get("sensor.my_station_water_level_stage").state
+        == STATE_UNAVAILABLE
+    )

tests/components/forked_daapd/test_media_player.py

@@ -345,12 +345,12 @@ async def mock_api_object_fixture(hass, config_entry, get_request_return_values)
 
 
 async def test_unload_config_entry(hass, config_entry, mock_api_object):
-    """Test the player is removed when the config entry is unloaded."""
+    """Test the player is set unavailable when the config entry is unloaded."""
     assert hass.states.get(TEST_MASTER_ENTITY_NAME)
     assert hass.states.get(TEST_ZONE_ENTITY_NAMES[0])
     await config_entry.async_unload(hass)
-    assert not hass.states.get(TEST_MASTER_ENTITY_NAME)
+    assert hass.states.get(TEST_MASTER_ENTITY_NAME).state == STATE_UNAVAILABLE
-    assert not hass.states.get(TEST_ZONE_ENTITY_NAMES[0])
+    assert hass.states.get(TEST_ZONE_ENTITY_NAMES[0]).state == STATE_UNAVAILABLE
 
 
 def test_master_state(hass, mock_api_object):

tests/components/fritzbox/test_init.py

@@ -4,7 +4,13 @@ from unittest.mock import Mock, call
 from homeassistant.components.fritzbox.const import DOMAIN as FB_DOMAIN
 from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
 from homeassistant.config_entries import ENTRY_STATE_LOADED, ENTRY_STATE_NOT_LOADED
-from homeassistant.const import CONF_DEVICES, CONF_HOST, CONF_PASSWORD, CONF_USERNAME
+from homeassistant.const import (
+    CONF_DEVICES,
+    CONF_HOST,
+    CONF_PASSWORD,
+    CONF_USERNAME,
+    STATE_UNAVAILABLE,
+)
 from homeassistant.helpers.typing import HomeAssistantType
 from homeassistant.setup import async_setup_component
 
@@ -45,8 +51,8 @@ async def test_setup_duplicate_config(hass: HomeAssistantType, fritz: Mock, capl
     assert "duplicate host entries found" in caplog.text
 
 
-async def test_unload(hass: HomeAssistantType, fritz: Mock):
+async def test_unload_remove(hass: HomeAssistantType, fritz: Mock):
-    """Test unload of integration."""
+    """Test unload and remove of integration."""
     fritz().get_devices.return_value = [FritzDeviceSwitchMock()]
     entity_id = f"{SWITCH_DOMAIN}.fake_name"
 
@@ -70,6 +76,14 @@ async def test_unload(hass: HomeAssistantType, fritz: Mock):
 
     await hass.config_entries.async_unload(entry.entry_id)
 
+    assert fritz().logout.call_count == 1
+    assert entry.state == ENTRY_STATE_NOT_LOADED
+    state = hass.states.get(entity_id)
+    assert state.state == STATE_UNAVAILABLE
+
+    await hass.config_entries.async_remove(entry.entry_id)
+    await hass.async_block_till_done()
+
     assert fritz().logout.call_count == 1
     assert entry.state == ENTRY_STATE_NOT_LOADED
     state = hass.states.get(entity_id)

tests/components/heos/test_media_player.py

@@ -587,10 +587,10 @@ async def test_select_input_command_error(
 
 
 async def test_unload_config_entry(hass, config_entry, config, controller):
-    """Test the player is removed when the config entry is unloaded."""
+    """Test the player is set unavailable when the config entry is unloaded."""
     await setup_platform(hass, config_entry, config)
     await config_entry.async_unload(hass)
-    assert not hass.states.get("media_player.test_player")
+    assert hass.states.get("media_player.test_player").state == STATE_UNAVAILABLE
 
 
 async def test_play_media_url(hass, config_entry, config, controller, caplog):

tests/components/homekit_controller/test_light.py

@@ -3,6 +3,7 @@ from aiohomekit.model.characteristics import CharacteristicsTypes
 from aiohomekit.model.services import ServicesTypes
 
 from homeassistant.components.homekit_controller.const import KNOWN_DEVICES
+from homeassistant.const import STATE_UNAVAILABLE
 
 from tests.components.homekit_controller.common import setup_test_component
 
@@ -209,8 +210,8 @@ async def test_light_becomes_unavailable_but_recovers(hass, utcnow):
     assert state.attributes["color_temp"] == 400
 
 
-async def test_light_unloaded(hass, utcnow):
+async def test_light_unloaded_removed(hass, utcnow):
-    """Test entity and HKDevice are correctly unloaded."""
+    """Test entity and HKDevice are correctly unloaded and removed."""
     helper = await setup_test_component(hass, create_lightbulb_service_with_color_temp)
 
     # Initial state is that the light is off
@@ -220,9 +221,15 @@ async def test_light_unloaded(hass, utcnow):
     unload_result = await helper.config_entry.async_unload(hass)
     assert unload_result is True
 
-    # Make sure entity is unloaded
+    # Make sure entity is set to unavailable state
-    assert hass.states.get(helper.entity_id) is None
+    assert hass.states.get(helper.entity_id).state == STATE_UNAVAILABLE
 
     # Make sure HKDevice is no longer set to poll this accessory
     conn = hass.data[KNOWN_DEVICES]["00:00:00:00:00:00"]
     assert not conn.pollable_characteristics
+
+    await helper.config_entry.async_remove(hass)
+    await hass.async_block_till_done()
+
+    # Make sure entity is removed
+    assert hass.states.get(helper.entity_id).state == STATE_UNAVAILABLE

tests/components/huisbaasje/test_init.py

@@ -11,7 +11,7 @@ from homeassistant.config_entries import (
     ENTRY_STATE_SETUP_ERROR,
     ConfigEntry,
 )
-from homeassistant.const import CONF_ID, CONF_PASSWORD, CONF_USERNAME
+from homeassistant.const import CONF_ID, CONF_PASSWORD, CONF_USERNAME, STATE_UNAVAILABLE
 from homeassistant.core import HomeAssistant
 from homeassistant.setup import async_setup_component
 
@@ -145,6 +145,14 @@ async def test_unload_entry(hass: HomeAssistant):
         await hass.config_entries.async_unload(config_entry.entry_id)
         assert config_entry.state == ENTRY_STATE_NOT_LOADED
         entities = hass.states.async_entity_ids("sensor")
+        assert len(entities) == 14
+        for entity in entities:
+            assert hass.states.get(entity).state == STATE_UNAVAILABLE
+
+        # Remove config entry
+        await hass.config_entries.async_remove(config_entry.entry_id)
+        await hass.async_block_till_done()
+        entities = hass.states.async_entity_ids("sensor")
         assert len(entities) == 0
 
         # Assert mocks are called

tests/components/input_boolean/test_init.py

@@ -264,7 +264,7 @@ async def test_reload(hass, hass_admin_user):
     assert "mdi:work_reloaded" == state_2.attributes.get(ATTR_ICON)
 
 
-async def test_load_person_storage(hass, storage_setup):
+async def test_load_from_storage(hass, storage_setup):
     """Test set up from storage."""
     assert await storage_setup()
     state = hass.states.get(f"{DOMAIN}.from_storage")

tests/components/met/test_weather.py

@@ -30,6 +30,7 @@ async def test_tracking_home(hass, mock_weather):
 
     entry = hass.config_entries.async_entries()[0]
     await hass.config_entries.async_remove(entry.entry_id)
+    await hass.async_block_till_done()
     assert len(hass.states.async_entity_ids("weather")) == 0
 
 
@@ -63,4 +64,5 @@ async def test_not_tracking_home(hass, mock_weather):
 
     entry = hass.config_entries.async_entries()[0]
     await hass.config_entries.async_remove(entry.entry_id)
+    await hass.async_block_till_done()
     assert len(hass.states.async_entity_ids("weather")) == 0

tests/components/nest/camera_sdm_test.py

@@ -339,6 +339,7 @@ async def test_camera_removed(hass, auth):
 
     for config_entry in hass.config_entries.async_entries(DOMAIN):
         await hass.config_entries.async_remove(config_entry.entry_id)
+    await hass.async_block_till_done()
     assert len(hass.states.async_all()) == 0
 
 

tests/components/nws/test_init.py

@@ -1,6 +1,7 @@
 """Tests for init module."""
 from homeassistant.components.nws.const import DOMAIN
 from homeassistant.components.weather import DOMAIN as WEATHER_DOMAIN
+from homeassistant.const import STATE_UNAVAILABLE
 
 from tests.common import MockConfigEntry
 from tests.components.nws.const import NWS_CONFIG
@@ -25,5 +26,12 @@ async def test_unload_entry(hass, mock_simple_nws):
     assert len(entries) == 1
 
     assert await hass.config_entries.async_unload(entries[0].entry_id)
-    assert len(hass.states.async_entity_ids(WEATHER_DOMAIN)) == 0
+    entities = hass.states.async_entity_ids(WEATHER_DOMAIN)
+    assert len(entities) == 1
+    for entity in entities:
+        assert hass.states.get(entity).state == STATE_UNAVAILABLE
     assert DOMAIN not in hass.data
+
+    assert await hass.config_entries.async_remove(entries[0].entry_id)
+    await hass.async_block_till_done()
+    assert len(hass.states.async_entity_ids(WEATHER_DOMAIN)) == 0

tests/components/ozw/test_init.py

@@ -4,6 +4,7 @@ from unittest.mock import patch
 from homeassistant import config_entries
 from homeassistant.components.hassio.handler import HassioAPIError
 from homeassistant.components.ozw import DOMAIN, PLATFORMS, const
+from homeassistant.const import ATTR_RESTORED, STATE_UNAVAILABLE
 
 from .common import setup_ozw
 
@@ -76,14 +77,21 @@ async def test_unload_entry(hass, generic_data, switch_msg, caplog):
     await hass.config_entries.async_unload(entry.entry_id)
 
     assert entry.state == config_entries.ENTRY_STATE_NOT_LOADED
-    assert len(hass.states.async_entity_ids("switch")) == 0
+    entities = hass.states.async_entity_ids("switch")
+    assert len(entities) == 1
+    for entity in entities:
+        assert hass.states.get(entity).state == STATE_UNAVAILABLE
+        assert hass.states.get(entity).attributes.get(ATTR_RESTORED)
 
     # Send a message for a switch from the broker to check that
     # all entity topic subscribers are unsubscribed.
     receive_message(switch_msg)
     await hass.async_block_till_done()
 
-    assert len(hass.states.async_entity_ids("switch")) == 0
+    assert len(hass.states.async_entity_ids("switch")) == 1
+    for entity in entities:
+        assert hass.states.get(entity).state == STATE_UNAVAILABLE
+        assert hass.states.get(entity).attributes.get(ATTR_RESTORED)
 
     # Load the integration again and check that there are no errors when
     # adding the entities.

tests/components/panasonic_viera/test_init.py

@@ -17,7 +17,7 @@ from homeassistant.components.panasonic_viera.const import (
     DOMAIN,
 )
 from homeassistant.config_entries import ENTRY_STATE_NOT_LOADED
-from homeassistant.const import CONF_HOST, CONF_NAME, CONF_PORT
+from homeassistant.const import CONF_HOST, CONF_NAME, CONF_PORT, STATE_UNAVAILABLE
 from homeassistant.setup import async_setup_component
 
 from tests.common import MockConfigEntry
@@ -253,9 +253,11 @@ async def test_setup_unload_entry(hass):
         await hass.async_block_till_done()
 
     await hass.config_entries.async_unload(mock_entry.entry_id)
-
     assert mock_entry.state == ENTRY_STATE_NOT_LOADED
-
     state = hass.states.get("media_player.panasonic_viera_tv")
+    assert state.state == STATE_UNAVAILABLE
 
+    await hass.config_entries.async_remove(mock_entry.entry_id)
+    await hass.async_block_till_done()
+    state = hass.states.get("media_player.panasonic_viera_tv")
     assert state is None

tests/components/plex/test_media_players.py

@@ -22,7 +22,7 @@ async def test_plex_tv_clients(
     media_players_after = len(hass.states.async_entity_ids("media_player"))
     assert media_players_after == media_players_before + 1
 
-    await hass.config_entries.async_unload(entry.entry_id)
+    await hass.config_entries.async_remove(entry.entry_id)
 
     # Ensure only plex.tv resource client is found
     with patch("plexapi.server.PlexServer.sessions", return_value=[]):

tests/components/smartthings/test_binary_sensor.py

@@ -12,7 +12,7 @@ from homeassistant.components.binary_sensor import (
 )
 from homeassistant.components.smartthings import binary_sensor
 from homeassistant.components.smartthings.const import DOMAIN, SIGNAL_SMARTTHINGS_UPDATE
-from homeassistant.const import ATTR_FRIENDLY_NAME
+from homeassistant.const import ATTR_FRIENDLY_NAME, STATE_UNAVAILABLE
 from homeassistant.helpers.dispatcher import async_dispatcher_send
 
 from .conftest import setup_platform
@@ -93,4 +93,7 @@ async def test_unload_config_entry(hass, device_factory):
     # Act
     await hass.config_entries.async_forward_entry_unload(config_entry, "binary_sensor")
     # Assert
-    assert not hass.states.get("binary_sensor.motion_sensor_1_motion")
+    assert (
+        hass.states.get("binary_sensor.motion_sensor_1_motion").state
+        == STATE_UNAVAILABLE
+    )

tests/components/smartthings/test_cover.py

@@ -19,7 +19,7 @@ from homeassistant.components.cover import (
     STATE_OPENING,
 )
 from homeassistant.components.smartthings.const import DOMAIN, SIGNAL_SMARTTHINGS_UPDATE
-from homeassistant.const import ATTR_BATTERY_LEVEL, ATTR_ENTITY_ID
+from homeassistant.const import ATTR_BATTERY_LEVEL, ATTR_ENTITY_ID, STATE_UNAVAILABLE
 from homeassistant.helpers.dispatcher import async_dispatcher_send
 
 from .conftest import setup_platform
@@ -193,4 +193,4 @@ async def test_unload_config_entry(hass, device_factory):
     # Act
     await hass.config_entries.async_forward_entry_unload(config_entry, COVER_DOMAIN)
     # Assert
-    assert not hass.states.get("cover.garage")
+    assert hass.states.get("cover.garage").state == STATE_UNAVAILABLE

tests/components/smartthings/test_fan.py

@@ -17,7 +17,11 @@ from homeassistant.components.fan import (
     SUPPORT_SET_SPEED,
 )
 from homeassistant.components.smartthings.const import DOMAIN, SIGNAL_SMARTTHINGS_UPDATE
-from homeassistant.const import ATTR_ENTITY_ID, ATTR_SUPPORTED_FEATURES
+from homeassistant.const import (
+    ATTR_ENTITY_ID,
+    ATTR_SUPPORTED_FEATURES,
+    STATE_UNAVAILABLE,
+)
 from homeassistant.helpers.dispatcher import async_dispatcher_send
 
 from .conftest import setup_platform
@@ -184,4 +188,4 @@ async def test_unload_config_entry(hass, device_factory):
     # Act
     await hass.config_entries.async_forward_entry_unload(config_entry, "fan")
     # Assert
-    assert not hass.states.get("fan.fan_1")
+    assert hass.states.get("fan.fan_1").state == STATE_UNAVAILABLE

tests/components/smartthings/test_light.py

@@ -19,7 +19,11 @@ from homeassistant.components.light import (
     SUPPORT_TRANSITION,
 )
 from homeassistant.components.smartthings.const import DOMAIN, SIGNAL_SMARTTHINGS_UPDATE
-from homeassistant.const import ATTR_ENTITY_ID, ATTR_SUPPORTED_FEATURES
+from homeassistant.const import (
+    ATTR_ENTITY_ID,
+    ATTR_SUPPORTED_FEATURES,
+    STATE_UNAVAILABLE,
+)
 from homeassistant.helpers.dispatcher import async_dispatcher_send
 
 from .conftest import setup_platform
@@ -304,4 +308,4 @@ async def test_unload_config_entry(hass, device_factory):
     # Act
     await hass.config_entries.async_forward_entry_unload(config_entry, "light")
     # Assert
-    assert not hass.states.get("light.color_dimmer_2")
+    assert hass.states.get("light.color_dimmer_2").state == STATE_UNAVAILABLE

tests/components/smartthings/test_lock.py

@@ -9,6 +9,7 @@ from pysmartthings.device import Status
 
 from homeassistant.components.lock import DOMAIN as LOCK_DOMAIN
 from homeassistant.components.smartthings.const import DOMAIN, SIGNAL_SMARTTHINGS_UPDATE
+from homeassistant.const import STATE_UNAVAILABLE
 from homeassistant.helpers.dispatcher import async_dispatcher_send
 
 from .conftest import setup_platform
@@ -104,4 +105,4 @@ async def test_unload_config_entry(hass, device_factory):
     # Act
     await hass.config_entries.async_forward_entry_unload(config_entry, "lock")
     # Assert
-    assert not hass.states.get("lock.lock_1")
+    assert hass.states.get("lock.lock_1").state == STATE_UNAVAILABLE

tests/components/smartthings/test_scene.py

@@ -5,7 +5,7 @@ The only mocking required is of the underlying SmartThings API object so
 real HTTP calls are not initiated during testing.
 """
 from homeassistant.components.scene import DOMAIN as SCENE_DOMAIN
-from homeassistant.const import ATTR_ENTITY_ID, SERVICE_TURN_ON
+from homeassistant.const import ATTR_ENTITY_ID, SERVICE_TURN_ON, STATE_UNAVAILABLE
 
 from .conftest import setup_platform
 
@@ -46,4 +46,4 @@ async def test_unload_config_entry(hass, scene):
     # Act
     await hass.config_entries.async_forward_entry_unload(config_entry, SCENE_DOMAIN)
     # Assert
-    assert not hass.states.get("scene.test_scene")
+    assert hass.states.get("scene.test_scene").state == STATE_UNAVAILABLE

tests/components/smartthings/test_sensor.py

@@ -13,6 +13,7 @@ from homeassistant.const import (
     ATTR_FRIENDLY_NAME,
     ATTR_UNIT_OF_MEASUREMENT,
     PERCENTAGE,
+    STATE_UNAVAILABLE,
     STATE_UNKNOWN,
 )
 from homeassistant.helpers.dispatcher import async_dispatcher_send
@@ -117,4 +118,4 @@ async def test_unload_config_entry(hass, device_factory):
     # Act
     await hass.config_entries.async_forward_entry_unload(config_entry, "sensor")
     # Assert
-    assert not hass.states.get("sensor.sensor_1_battery")
+    assert hass.states.get("sensor.sensor_1_battery").state == STATE_UNAVAILABLE

tests/components/smartthings/test_switch.py

@@ -12,6 +12,7 @@ from homeassistant.components.switch import (
     ATTR_TODAY_ENERGY_KWH,
     DOMAIN as SWITCH_DOMAIN,
 )
+from homeassistant.const import STATE_UNAVAILABLE
 from homeassistant.helpers.dispatcher import async_dispatcher_send
 
 from .conftest import setup_platform
@@ -96,4 +97,4 @@ async def test_unload_config_entry(hass, device_factory):
     # Act
     await hass.config_entries.async_forward_entry_unload(config_entry, "switch")
     # Assert
-    assert not hass.states.get("switch.switch_1")
+    assert hass.states.get("switch.switch_1").state == STATE_UNAVAILABLE

tests/components/vizio/test_init.py

@@ -3,6 +3,7 @@ import pytest
 
 from homeassistant.components.media_player.const import DOMAIN as MP_DOMAIN
 from homeassistant.components.vizio.const import DOMAIN
+from homeassistant.const import STATE_UNAVAILABLE
 from homeassistant.helpers.typing import HomeAssistantType
 from homeassistant.setup import async_setup_component
 
@@ -41,7 +42,10 @@ async def test_tv_load_and_unload(
 
     assert await config_entry.async_unload(hass)
     await hass.async_block_till_done()
-    assert len(hass.states.async_entity_ids(MP_DOMAIN)) == 0
+    entities = hass.states.async_entity_ids(MP_DOMAIN)
+    assert len(entities) == 1
+    for entity in entities:
+        assert hass.states.get(entity).state == STATE_UNAVAILABLE
     assert DOMAIN not in hass.data
 
 
@@ -62,5 +66,8 @@ async def test_speaker_load_and_unload(
 
     assert await config_entry.async_unload(hass)
     await hass.async_block_till_done()
-    assert len(hass.states.async_entity_ids(MP_DOMAIN)) == 0
+    entities = hass.states.async_entity_ids(MP_DOMAIN)
+    assert len(entities) == 1
+    for entity in entities:
+        assert hass.states.get(entity).state == STATE_UNAVAILABLE
     assert DOMAIN not in hass.data

tests/components/yeelight/test_init.py

@@ -11,7 +11,7 @@ from homeassistant.components.yeelight import (
     DOMAIN,
     NIGHTLIGHT_SWITCH_TYPE_LIGHT,
 )
-from homeassistant.const import CONF_DEVICES, CONF_HOST, CONF_NAME
+from homeassistant.const import CONF_DEVICES, CONF_HOST, CONF_NAME, STATE_UNAVAILABLE
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry
 from homeassistant.setup import async_setup_component
@@ -50,6 +50,12 @@ async def test_setup_discovery(hass: HomeAssistant):
 
     # Unload
     assert await hass.config_entries.async_unload(config_entry.entry_id)
+    assert hass.states.get(ENTITY_BINARY_SENSOR).state == STATE_UNAVAILABLE
+    assert hass.states.get(ENTITY_LIGHT).state == STATE_UNAVAILABLE
+
+    # Remove
+    assert await hass.config_entries.async_remove(config_entry.entry_id)
+    await hass.async_block_till_done()
     assert hass.states.get(ENTITY_BINARY_SENSOR) is None
     assert hass.states.get(ENTITY_LIGHT) is None
 

tests/helpers/test_collection.py

@@ -226,7 +226,7 @@ async def test_attach_entity_component_collection(hass):
     """Test attaching collection to entity component."""
     ent_comp = entity_component.EntityComponent(_LOGGER, "test", hass)
     coll = collection.ObservableCollection(_LOGGER)
-    collection.attach_entity_component_collection(ent_comp, coll, MockEntity)
+    collection.sync_entity_lifecycle(hass, "test", "test", ent_comp, coll, MockEntity)
 
     await coll.notify_changes(
         [

tests/helpers/test_entity.py

@@ -7,7 +7,7 @@ from unittest.mock import MagicMock, PropertyMock, patch
 
 import pytest
 
-from homeassistant.const import ATTR_DEVICE_CLASS, STATE_UNAVAILABLE
+from homeassistant.const import ATTR_DEVICE_CLASS, STATE_UNAVAILABLE, STATE_UNKNOWN
 from homeassistant.core import Context
 from homeassistant.helpers import entity, entity_registry
 
@@ -718,3 +718,29 @@ async def test_setup_source(hass):
     await platform.async_reset()
 
     assert entity.entity_sources(hass) == {}
+
+
+async def test_removing_entity_unavailable(hass):
+    """Test removing an entity that is still registered creates an unavailable state."""
+    entry = entity_registry.RegistryEntry(
+        entity_id="hello.world",
+        unique_id="test-unique-id",
+        platform="test-platform",
+        disabled_by=None,
+    )
+
+    ent = entity.Entity()
+    ent.hass = hass
+    ent.entity_id = "hello.world"
+    ent.registry_entry = entry
+    ent.async_write_ha_state()
+
+    state = hass.states.get("hello.world")
+    assert state is not None
+    assert state.state == STATE_UNKNOWN
+
+    await ent.async_remove()
+
+    state = hass.states.get("hello.world")
+    assert state is not None
+    assert state.state == STATE_UNAVAILABLE