Reload config entry when entity enabled in entity registry, remove entity if disabled. (#26120)

* Reload config entry when disabled_by updated in entity registry

* Add types

* Remove entities that get disabled

* Remove unnecessary domain checks.

* Attach handler in async_setup

* Remove unused var

* Type

* Fix test

* Fix tests

tests/test_config_entries.py

@@ -20,6 +20,7 @@ from tests.common import (
     MockEntity,
     mock_integration,
     mock_entity_platform,
+    mock_registry,
 )
 
 
@@ -925,3 +926,78 @@ async def test_init_custom_integration(hass):
             return_value=mock_coro(integration),
         ):
             await hass.config_entries.flow.async_init("bla")
+
+
+async def test_support_entry_unload(hass):
+    """Test unloading entry."""
+    assert await config_entries.support_entry_unload(hass, "light")
+    assert not await config_entries.support_entry_unload(hass, "auth")
+
+
+async def test_reload_entry_entity_registry_ignores_no_entry(hass):
+    """Test reloading entry in entity registry skips if no config entry linked."""
+    handler = config_entries.EntityRegistryDisabledHandler(hass)
+    registry = mock_registry(hass)
+
+    # Test we ignore entities without config entry
+    entry = registry.async_get_or_create("light", "hue", "123")
+    registry.async_update_entity(entry.entity_id, disabled_by="user")
+    await hass.async_block_till_done()
+    assert not handler.changed
+    assert handler._remove_call_later is None
+
+
+async def test_reload_entry_entity_registry_works(hass):
+    """Test we schedule an entry to be reloaded if disabled_by is updated."""
+    handler = config_entries.EntityRegistryDisabledHandler(hass)
+    handler.async_setup()
+    registry = mock_registry(hass)
+
+    config_entry = MockConfigEntry(
+        domain="comp", state=config_entries.ENTRY_STATE_LOADED
+    )
+    config_entry.add_to_hass(hass)
+    mock_setup_entry = MagicMock(return_value=mock_coro(True))
+    mock_unload_entry = MagicMock(return_value=mock_coro(True))
+    mock_integration(
+        hass,
+        MockModule(
+            "comp",
+            async_setup_entry=mock_setup_entry,
+            async_unload_entry=mock_unload_entry,
+        ),
+    )
+    mock_entity_platform(hass, "config_flow.comp", None)
+
+    # Only changing disabled_by should update trigger
+    entity_entry = registry.async_get_or_create(
+        "light", "hue", "123", config_entry=config_entry
+    )
+    registry.async_update_entity(entity_entry.entity_id, name="yo")
+    await hass.async_block_till_done()
+    assert not handler.changed
+    assert handler._remove_call_later is None
+
+    # Disable entity, we should not do anything, only act when enabled.
+    registry.async_update_entity(entity_entry.entity_id, disabled_by="user")
+    await hass.async_block_till_done()
+    assert not handler.changed
+    assert handler._remove_call_later is None
+
+    # Enable entity, check we are reloading config entry.
+    registry.async_update_entity(entity_entry.entity_id, disabled_by=None)
+    await hass.async_block_till_done()
+    assert handler.changed == {config_entry.entry_id}
+    assert handler._remove_call_later is not None
+
+    async_fire_time_changed(
+        hass,
+        dt.utcnow()
+        + timedelta(
+            seconds=config_entries.EntityRegistryDisabledHandler.RELOAD_AFTER_UPDATE_DELAY
+            + 1
+        ),
+    )
+    await hass.async_block_till_done()
+
+    assert len(mock_unload_entry.mock_calls) == 1