From 5f22b0c6cc50b08a58d9bf1948b90c03baf50098 Mon Sep 17 00:00:00 2001
From: puddly <32534428+puddly@users.noreply.github.com>
Date: Thu, 4 May 2023 20:02:17 -0400
Subject: [PATCH] Handle invalid ZHA cluster handlers (#92543)

* Do not crash on startup when an invalid cluster handler is encountered

* Add a unit test
---
 homeassistant/components/zha/core/endpoint.py | 14 ++++++-
 tests/components/zha/test_cluster_handlers.py | 40 +++++++++++++++++++
 2 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/homeassistant/components/zha/core/endpoint.py b/homeassistant/components/zha/core/endpoint.py
index d134c033ed7..53a3fb883ef 100644
--- a/homeassistant/components/zha/core/endpoint.py
+++ b/homeassistant/components/zha/core/endpoint.py
@@ -137,7 +137,19 @@ class Endpoint:
             ):
                 cluster_handler_class = MultistateInput
             # end of ugly hack
-            cluster_handler = cluster_handler_class(cluster, self)
+
+            try:
+                cluster_handler = cluster_handler_class(cluster, self)
+            except KeyError as err:
+                _LOGGER.warning(
+                    "Cluster handler %s for cluster %s on endpoint %s is invalid: %s",
+                    cluster_handler_class,
+                    cluster,
+                    self,
+                    err,
+                )
+                continue
+
             if cluster_handler.name == const.CLUSTER_HANDLER_POWER_CONFIGURATION:
                 self._device.power_configuration_ch = cluster_handler
             elif cluster_handler.name == const.CLUSTER_HANDLER_IDENTIFY:
diff --git a/tests/components/zha/test_cluster_handlers.py b/tests/components/zha/test_cluster_handlers.py
index c0c455542d3..1897383b6c4 100644
--- a/tests/components/zha/test_cluster_handlers.py
+++ b/tests/components/zha/test_cluster_handlers.py
@@ -1,11 +1,13 @@
 """Test ZHA Core cluster handlers."""
 import asyncio
 from collections.abc import Callable
+import logging
 import math
 from unittest import mock
 from unittest.mock import AsyncMock, patch
 
 import pytest
+import zigpy.device
 import zigpy.endpoint
 from zigpy.endpoint import Endpoint as ZigpyEndpoint
 import zigpy.profiles.zha
@@ -791,3 +793,41 @@ async def test_configure_reporting(hass: HomeAssistant, endpoint) -> None:
             }
         ),
     ]
+
+
+async def test_invalid_cluster_handler(hass: HomeAssistant, caplog) -> None:
+    """Test setting up a cluster handler that fails to match properly."""
+
+    class TestZigbeeClusterHandler(cluster_handlers.ClusterHandler):
+        REPORT_CONFIG = (
+            cluster_handlers.AttrReportConfig(attr="missing_attr", config=(1, 60, 1)),
+        )
+
+    mock_device = mock.AsyncMock(spec_set=zigpy.device.Device)
+    zigpy_ep = zigpy.endpoint.Endpoint(mock_device, endpoint_id=1)
+
+    cluster = zigpy_ep.add_input_cluster(zigpy.zcl.clusters.lighting.Color.cluster_id)
+    cluster.configure_reporting_multiple = AsyncMock(
+        spec_set=cluster.configure_reporting_multiple,
+        return_value=[
+            foundation.ConfigureReportingResponseRecord(
+                status=foundation.Status.SUCCESS
+            )
+        ],
+    )
+
+    mock_zha_device = mock.AsyncMock(spec_set=ZHADevice)
+    zha_endpoint = Endpoint(zigpy_ep, mock_zha_device)
+
+    # The cluster handler throws an error when matching this cluster
+    with pytest.raises(KeyError):
+        TestZigbeeClusterHandler(cluster, zha_endpoint)
+
+    # And one is also logged at runtime
+    with patch.dict(
+        registries.ZIGBEE_CLUSTER_HANDLER_REGISTRY,
+        {cluster.cluster_id: TestZigbeeClusterHandler},
+    ), caplog.at_level(logging.WARNING):
+        zha_endpoint.add_all_cluster_handlers()
+
+    assert "missing_attr" in caplog.text