Fix ZHA update entities not working after reload (#164290)

homeassistant/components/zha/__init__.py

@@ -274,6 +274,9 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
 
 async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
     """Unload ZHA config entry."""
+    if not await hass.config_entries.async_unload_platforms(config_entry, PLATFORMS):
+        return False
+
     ha_zha_data = get_zha_data(hass)
     ha_zha_data.config_entry = None
 
@@ -281,6 +284,8 @@ async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) ->
         await ha_zha_data.gateway_proxy.shutdown()
         ha_zha_data.gateway_proxy = None
 
+    ha_zha_data.update_coordinator = None
+
     # clean up any remaining entity metadata
     # (entities that have been discovered but not yet added to HA)
     # suppress KeyError because we don't know what state we may
@@ -291,7 +296,7 @@ async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) ->
 
     websocket_api.async_unload_api(hass)
 
-    return await hass.config_entries.async_unload_platforms(config_entry, PLATFORMS)
+    return True
 
 
 async def async_migrate_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:

tests/components/zha/test_update.py

@@ -33,6 +33,7 @@ from homeassistant.components.update import (
 from homeassistant.components.zha.helpers import (
     ZHADeviceProxy,
     ZHAGatewayProxy,
+    get_zha_data,
     get_zha_gateway,
     get_zha_gateway_proxy,
 )
@@ -55,6 +56,7 @@ from homeassistant.setup import async_setup_component
 from .common import find_entity_id, update_attribute_cache
 from .conftest import SIG_EP_INPUT, SIG_EP_OUTPUT, SIG_EP_PROFILE, SIG_EP_TYPE
 
+from tests.common import MockConfigEntry
 from tests.typing import WebSocketGenerator
 
 
@@ -267,6 +269,43 @@ async def test_firmware_update_notification_from_service_call(
         )
 
 
+async def test_firmware_update_poll_after_reload(
+    hass: HomeAssistant,
+    setup_zha: Callable[..., Coroutine[None]],
+    config_entry: MockConfigEntry,
+    zigpy_device_mock: Callable[..., Device],
+) -> None:
+    """Test polling a ZHA update entity still works after reloading ZHA."""
+    await setup_zha()
+    await async_setup_component(hass, HA_DOMAIN, {})
+
+    zha_data = get_zha_data(hass)
+    coordinator_before = zha_data.update_coordinator
+    assert coordinator_before is not None
+
+    assert await hass.config_entries.async_reload(config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    coordinator_after = get_zha_data(hass).update_coordinator
+    assert coordinator_after is not None
+    assert coordinator_after is not coordinator_before
+
+    zha_device, _, _, _ = await setup_test_data(hass, zigpy_device_mock)
+    entity_id = find_entity_id(Platform.UPDATE, zha_device, hass)
+    assert entity_id is not None
+
+    with patch("zigpy.ota.OTA.broadcast_notify") as mock_broadcast_notify:
+        await hass.services.async_call(
+            HA_DOMAIN,
+            SERVICE_UPDATE_ENTITY,
+            service_data={ATTR_ENTITY_ID: entity_id},
+            blocking=True,
+        )
+
+        assert mock_broadcast_notify.await_count == 1
+        assert mock_broadcast_notify.call_args_list[0] == call(jitter=100)
+
+
 def make_packet(zigpy_device, cluster, cmd_name: str, **kwargs):
     """Make a zigpy packet."""
     req_hdr, req_cmd = cluster._create_request(